From 46f10478dc8f1a4dcc150b424dc6210b0c861dd9 Mon Sep 17 00:00:00 2001 From: Kotaro Kikuchi Date: Sat, 4 Jul 2026 12:06:24 +0900 Subject: [PATCH 1/4] [Bugfix] Qwen-Image-Layered: allow layers=2 to match upstream demo (#4777) Signed-off-by: Kotaro Kikuchi --- .../online_serving/test_qwen_image_layered_expansion.py | 6 +++--- tests/entrypoints/openai_api/test_image_server.py | 8 ++++---- vllm_omni/entrypoints/openai/image_api_utils.py | 2 +- vllm_omni/entrypoints/openai/protocol/images.py | 2 +- 4 files changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/e2e/online_serving/test_qwen_image_layered_expansion.py b/tests/e2e/online_serving/test_qwen_image_layered_expansion.py index c26e21039..315fba6ef 100644 --- a/tests/e2e/online_serving/test_qwen_image_layered_expansion.py +++ b/tests/e2e/online_serving/test_qwen_image_layered_expansion.py @@ -7,7 +7,7 @@ and are supported by Qwen-Image-Layered model. Kept cases (maximum feature coverage, no cpu-offload): sp_001 : cache_dit + Ulysses-SP 2 (2×H100) cfg_parallel_002 : cache_dit + CFG-Parallel 2 (2×H100) - layers_guard_001 : layerwise offload + layers=3 (1×H100, issue #1969 guard) + layers_guard_001 : layerwise offload + layers=2 (1×H100, issue #1969 guard) Total distinct features covered: cache_dit, Ulysses-SP, CFG-Parallel, layerwise-offload. """ @@ -116,8 +116,8 @@ LAYERS_GUARD_CASES = [ model=MODEL, server_args=["--enable-layerwise-offload"], ), - 3, - id="layers_guard_001_layers3", + 2, + id="layers_guard_001_layers2", marks=SINGLE_CARD_FEATURE_MARKS, ), ] diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 31a82102a..5a7ef7bf0 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -1244,7 +1244,7 @@ def test_parameter_validation(): # Invalid layers for layered models (must stay within the backend-supported range) with pytest.raises(ValueError): - ImageGenerationRequest(prompt="test", layers=2) + ImageGenerationRequest(prompt="test", layers=1) with pytest.raises(ValueError): ImageGenerationRequest(prompt="test", layers=11) @@ -1663,13 +1663,13 @@ def test_image_edit_invalid_layers(async_omni_test_client): files=[("image", img_bytes)], data={ "prompt": "test", - "layers": 2, + "layers": 1, }, ) assert response.status_code == 400 detail = response.json()["detail"] assert "Invalid layers" in detail - assert "layers must be between 3 and 10 inclusive" in detail + assert "layers must be between 2 and 10 inclusive" in detail # Test layers above the supported range response = async_omni_test_client.post( @@ -1683,7 +1683,7 @@ def test_image_edit_invalid_layers(async_omni_test_client): assert response.status_code == 400 detail = response.json()["detail"] assert "Invalid layers" in detail - assert "layers must be between 3 and 10 inclusive" in detail + assert "layers must be between 2 and 10 inclusive" in detail def test_image_edit_resolution_and_size_conflict(async_omni_test_client): diff --git a/vllm_omni/entrypoints/openai/image_api_utils.py b/vllm_omni/entrypoints/openai/image_api_utils.py index b958d786a..5bd79e936 100644 --- a/vllm_omni/entrypoints/openai/image_api_utils.py +++ b/vllm_omni/entrypoints/openai/image_api_utils.py @@ -14,7 +14,7 @@ import io import PIL.Image SUPPORTED_LAYERED_RESOLUTIONS = (640, 1024) -SUPPORTED_LAYERED_LAYERS_RANGE = range(3, 11) +SUPPORTED_LAYERED_LAYERS_RANGE = range(2, 11) def parse_size(size_str: str) -> tuple[int, int]: diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py index 78f63d387..b7247962c 100644 --- a/vllm_omni/entrypoints/openai/protocol/images.py +++ b/vllm_omni/entrypoints/openai/protocol/images.py @@ -60,7 +60,7 @@ class ImageGenerationRequest(BaseModel): user: str | None = Field(default=None, description="User identifier for tracking") layers: int | None = Field( default=None, - description="Number of output layers for layered image models. Supported range: 3-10.", + description="Number of output layers for layered image models. Supported range: 2-10.", ) @field_validator("size") From f7cdc45977958cd8901edccb11d8f481bca97f0c Mon Sep 17 00:00:00 2001 From: Haco <923390377@qq.com> Date: Sat, 4 Jul 2026 11:17:25 +0800 Subject: [PATCH 2/4] [BugFix] deploy schema stage runtime overrides for CLI diffusion flags no longer works (#4101) Signed-off-by: xiaohajiayou <923390377@qq.com> --- tests/test_config_factory.py | 191 +++++++++++++++++++++++++++++++ vllm_omni/config/stage_config.py | 87 +++++++++++--- 2 files changed, 262 insertions(+), 16 deletions(-) diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index 63a67d652..23d36c338 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -1588,6 +1588,197 @@ class TestAuraOmniDeploy: assert stages[3].yaml_engine_args["model_arch"] == "Qwen3TTSCode2Wav" +class TestDeployCliOverrideFlow: + """Test deploy-YAML baselines overridden by CLI runtime overrides.""" + + def test_diffusion_deploy_fields_can_be_overridden_by_cli(self, tmp_path): + deploy_path = tmp_path / "diffusion_stage.yaml" + deploy_path.write_text( + """ +async_chunk: false +stages: + - stage_id: 0 + devices: "4,5" + parallel_config: + pipeline_parallel_size: 1 + data_parallel_size: 1 + tensor_parallel_size: 1 + enable_expert_parallel: false + sequence_parallel_size: 1 + ulysses_degree: 1 + ring_degree: 1 + cfg_parallel_size: 1 + vae_patch_parallel_size: 1 + use_hsdp: false + hsdp_shard_size: -1 + hsdp_replicate_size: 1 + engine_args: + cache_backend: cache_dit + diffusion_attention_backend: FLASH_ATTN + diffusion_kv_cache_dtype: auto + step_execution: false + vae_use_tiling: false + enable_cpu_offload: false + max_generated_image_size: 1048576 + tts_max_instructions_length: 1000 +""", + encoding="utf-8", + ) + + pipeline = PipelineConfig( + model_type="test_diffusion", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + final_output=True, + final_output_type="image", + ), + ), + ) + deploy = load_deploy_config(deploy_path) + stages = merge_pipeline_deploy(pipeline, deploy) + stage = stages[0] + + assert stage.yaml_engine_args["parallel_config"]["pipeline_parallel_size"] == 1 + assert stage.yaml_engine_args["parallel_config"]["data_parallel_size"] == 1 + assert stage.yaml_engine_args["parallel_config"]["tensor_parallel_size"] == 1 + assert stage.yaml_engine_args["parallel_config"]["enable_expert_parallel"] is False + assert stage.yaml_engine_args["parallel_config"]["sequence_parallel_size"] == 1 + assert stage.yaml_engine_args["parallel_config"]["ulysses_degree"] == 1 + assert stage.yaml_engine_args["parallel_config"]["ring_degree"] == 1 + assert stage.yaml_engine_args["parallel_config"]["cfg_parallel_size"] == 1 + assert stage.yaml_engine_args["parallel_config"]["vae_patch_parallel_size"] == 1 + assert stage.yaml_engine_args["cache_backend"] == "cache_dit" + assert stage.yaml_engine_args["parallel_config"]["use_hsdp"] is False + assert stage.yaml_engine_args["parallel_config"]["hsdp_shard_size"] == -1 + assert stage.yaml_engine_args["parallel_config"]["hsdp_replicate_size"] == 1 + assert stage.yaml_engine_args["diffusion_attention_backend"] == "FLASH_ATTN" + assert stage.yaml_engine_args["diffusion_kv_cache_dtype"] == "auto" + assert stage.yaml_engine_args["step_execution"] is False + assert stage.yaml_engine_args["vae_use_tiling"] is False + assert stage.yaml_engine_args["enable_cpu_offload"] is False + assert stage.yaml_engine_args["max_generated_image_size"] == 1048576 + assert stage.yaml_engine_args["tts_max_instructions_length"] == 1000 + + stage.runtime_overrides = StageConfigFactory._merge_cli_overrides( + stage, + { + "pipeline_parallel_size": 2, + "data_parallel_size": 3, + "tensor_parallel_size": 4, + "enable_expert_parallel": True, + "sequence_parallel_size": 24, + "ulysses_degree": 2, + "ring_degree": 4, + "cfg_parallel_size": 2, + "vae_patch_parallel_size": 2, + "cache_backend": "tea_cache", + "use_hsdp": True, + "hsdp_shard_size": 8, + "hsdp_replicate_size": 2, + "diffusion_attention_backend": "SAGE_ATTN", + "diffusion_kv_cache_dtype": "fp8", + "step_execution": True, + "vae_use_tiling": True, + "enable_cpu_offload": True, + "max_generated_image_size": 2097152, + "tts_max_instructions_length": 2000, + }, + ) + + omega_config = stage.to_omegaconf() + + assert omega_config.engine_args.cache_backend == "tea_cache" + assert omega_config.engine_args.diffusion_attention_backend == "SAGE_ATTN" + assert omega_config.engine_args.diffusion_kv_cache_dtype == "fp8" + assert omega_config.engine_args.step_execution is True + assert omega_config.engine_args.vae_use_tiling is True + assert omega_config.engine_args.enable_cpu_offload is True + assert omega_config.engine_args.max_generated_image_size == 2097152 + assert omega_config.engine_args.tts_max_instructions_length == 2000 + assert omega_config.engine_args.parallel_config.pipeline_parallel_size == 2 + assert omega_config.engine_args.parallel_config.data_parallel_size == 3 + assert omega_config.engine_args.parallel_config.tensor_parallel_size == 4 + assert omega_config.engine_args.parallel_config.enable_expert_parallel is True + assert omega_config.engine_args.parallel_config.sequence_parallel_size == 24 + assert omega_config.engine_args.parallel_config.ulysses_degree == 2 + assert omega_config.engine_args.parallel_config.ring_degree == 4 + assert omega_config.engine_args.parallel_config.cfg_parallel_size == 2 + assert omega_config.engine_args.parallel_config.vae_patch_parallel_size == 2 + assert omega_config.engine_args.parallel_config.use_hsdp is True + assert omega_config.engine_args.parallel_config.hsdp_shard_size == 8 + assert omega_config.engine_args.parallel_config.hsdp_replicate_size == 2 + + def test_llm_deploy_fields_can_be_overridden_by_cli(self, tmp_path): + deploy_path = tmp_path / "llm_stage.yaml" + deploy_path.write_text( + """ +async_chunk: false +stages: + - stage_id: 0 + devices: "0,1" + engine_args: + tensor_parallel_size: 1 + enable_expert_parallel: false + gpu_memory_utilization: 0.5 + max_num_seqs: 16 + max_num_batched_tokens: 1024 + max_model_len: 4096 + enforce_eager: false +""", + encoding="utf-8", + ) + + pipeline = PipelineConfig( + model_type="test_llm", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="thinker", + execution_type=StageExecutionType.LLM_AR, + final_output=True, + final_output_type="text", + ), + ), + ) + deploy = load_deploy_config(deploy_path) + stages = merge_pipeline_deploy(pipeline, deploy) + stage = stages[0] + + assert stage.yaml_engine_args["tensor_parallel_size"] == 1 + assert stage.yaml_engine_args["enable_expert_parallel"] is False + assert stage.yaml_engine_args["gpu_memory_utilization"] == 0.5 + assert stage.yaml_engine_args["max_num_seqs"] == 16 + assert stage.yaml_engine_args["max_num_batched_tokens"] == 1024 + assert stage.yaml_engine_args["max_model_len"] == 4096 + assert stage.yaml_engine_args["enforce_eager"] is False + + stage.runtime_overrides = StageConfigFactory._merge_cli_overrides( + stage, + { + "tensor_parallel_size": 2, + "enable_expert_parallel": True, + "gpu_memory_utilization": 0.9, + "max_num_seqs": 32, + "max_num_batched_tokens": 2048, + "max_model_len": 8192, + "enforce_eager": True, + }, + ) + + omega_config = stage.to_omegaconf() + + assert omega_config.engine_args.tensor_parallel_size == 2 + assert omega_config.engine_args.enable_expert_parallel is True + assert omega_config.engine_args.gpu_memory_utilization == 0.9 + assert omega_config.engine_args.max_num_seqs == 32 + assert omega_config.engine_args.max_num_batched_tokens == 2048 + assert omega_config.engine_args.max_model_len == 8192 + assert omega_config.engine_args.enforce_eager is True + + class TestSentinelDefaultPrecedence: """Caller-typed (non-None) values win over YAML; None values fall through to YAML / dataclass defaults (#3035).""" diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index a269216da..f66023bc4 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -63,7 +63,12 @@ def build_stage_runtime_overrides( if internal_keys is None: from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys - internal_keys = internal_blacklist_keys() | SHARED_FIELDS + # Some fields are modeled as orchestrator-owned for top-level CLI + # parsing, but are also legitimate deploy-time stage overrides. Keep + # the default blacklist for true orchestrator/shared fields while + # allowing any field explicitly represented by the deploy schema to + # continue flowing into per-stage overrides. + internal_keys = (internal_blacklist_keys() | SHARED_FIELDS) - deploy_runtime_override_keys() result: dict[str, Any] = {} @@ -303,7 +308,7 @@ class StageDeployConfig: the top level of ``DeployConfig`` and propagated to every stage. """ - # === Omni fields === + # === Omni stage wrapper fields === # Stage identity and Omni runtime placement. stage_id: int devices: str | None = None @@ -316,22 +321,33 @@ class StageDeployConfig: default_sampling_params: dict[str, Any] | None = None subtalker_sampling_params: dict[str, Any] | None = None - # === vLLM EngineArgs fields === - # Parallelism and scheduler/memory capacity. + # === Generic stage engine fields === + # Parallelism, scheduler, and memory-capacity controls. tensor_parallel_size: int | None = None + enable_expert_parallel: bool | None = None gpu_memory_utilization: float | None = None max_num_seqs: int | None = None max_num_batched_tokens: int | None = None max_model_len: int | None = None - # Execution, scheduling, and KV/cache behavior. + # Generic execution, scheduling, and KV/cache behavior. enforce_eager: bool | None = None async_scheduling: bool | None = None disable_hybrid_kv_cache_manager: bool | None = None mm_processor_cache_gb: float | None = None - # Diffusion parallel_config deploy override fields. - enable_expert_parallel: bool | None = None + # Generic compilation, profiling, tokenizer/config parsing, and model + # loading controls. + compilation_config: dict[str, Any] | None = None + profiler_config: dict[str, Any] | None = None + skip_mm_profiling: bool | None = None + enable_flashinfer_autotune: bool | None = None + config_format: str | None = None + load_format: str | None = None + tokenizer_mode: str | None = None + + # === Diffusion stage runtime fields === + # Diffusion parallel_config deploy/runtime override fields. ulysses_degree: int | None = None ulysses_mode: str | None = None ring_degree: int | None = None @@ -343,16 +359,44 @@ class StageDeployConfig: hsdp_shard_size: int | None = None hsdp_replicate_size: int | None = None - # Compilation, profiling, tokenizer/config parsing, and model loading. - compilation_config: dict[str, Any] | None = None - profiler_config: dict[str, Any] | None = None - skip_mm_profiling: bool | None = None - enable_flashinfer_autotune: bool | None = None - config_format: str | None = None - load_format: str | None = None - tokenizer_mode: str | None = None + # Diffusion model loading and adapter construction. + model_class_name: str | None = None + diffusion_load_format: str | None = None + diffusers_load_kwargs: dict[str, Any] | None = None + diffusers_call_kwargs: dict[str, Any] | None = None + diffusion_quantization_config: str | None = None + diffusion_attention_backend: str | None = None + diffusion_attention_config: dict[str, Any] | None = None - # Pass-through vLLM EngineArgs fields that are not represented above. + # Diffusion execution, cache, and VAE behavior. + cache_backend: str | None = None + cache_config: dict[str, Any] | None = None + enable_cache_dit_summary: bool | None = None + step_execution: bool | None = None + vae_use_slicing: bool | None = None + vae_use_tiling: bool | None = None + boundary_ratio: float | None = None + flow_shift: float | None = None + diffusion_kv_cache_dtype: str | None = None + diffusion_kv_cache_skip_steps: str | None = None + diffusion_kv_cache_skip_layers: str | None = None + auxiliary_text_encoder: str | None = None + + # Runtime optimizations used by diffusion loading/execution. + enable_multithread_weight_load: bool | None = None + num_weight_load_threads: int | None = None + enable_cpu_offload: bool | None = None + enable_layerwise_offload: bool | None = None + + # Diffusion-specific debug and observability knobs. + enable_diffusion_pipeline_profiler: bool | None = None + + # Modality/service constraints consumed outside the core engine config. + max_generated_image_size: int | None = None + tts_max_instructions_length: int | None = None + + # === Pass-through stage engine fields === + # Pass-through stage engine args that are not represented above. engine_extras: dict[str, Any] = field(default_factory=dict) @@ -409,6 +453,17 @@ _STAGE_RESERVED_KEYS = frozenset( _STAGE_DEPLOY_FIELDS = {f.name: f for f in fields(StageDeployConfig) if f.name not in _STAGE_RESERVED_KEYS} +def deploy_runtime_override_keys() -> frozenset[str]: + """Return deploy-schema fields that are valid CLI/runtime overrides. + + These keys form the positive contract for stage override propagation: + stage-scoped deploy knobs plus top-level pipeline-wide engine settings. + They must remain overridable even if they are also modeled on + ``OrchestratorArgs`` for top-level CLI parsing. + """ + return frozenset(_STAGE_DEPLOY_FIELDS) | frozenset(_PIPELINE_WIDE_ENGINE_FIELDS) + + def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig: """Parse a single stage entry from deploy YAML into StageDeployConfig.""" # Get the non-reserved keys for this stage From e73245c60fca55eba5e86b2704e8e266e88ad93b Mon Sep 17 00:00:00 2001 From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> Date: Sat, 4 Jul 2026 11:45:49 +0800 Subject: [PATCH 3/4] [Bugfix] Adapt ZImagePipeline.forward calling in MingImagePipeline with DiffusionRequestBatch (#4866) Signed-off-by: Yuanheng Zhao --- .../ming_flash_omni/pipeline_ming_imagegen.py | 89 ++++++++----------- 1 file changed, 39 insertions(+), 50 deletions(-) diff --git a/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py b/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py index f3e0f3357..c67f8cb24 100644 --- a/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py +++ b/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py @@ -24,7 +24,6 @@ from __future__ import annotations import logging import os -from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -49,7 +48,9 @@ from vllm_omni.diffusion.models.ming_flash_omni.ming_zimage_transformer import ( MingZImageTransformer2DModel, ) from vllm_omni.diffusion.models.z_image.pipeline_z_image import ZImagePipeline +from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) @@ -58,38 +59,6 @@ from vllm_omni.transformers_utils.configs.ming_flash_omni import MingImageGenCon logger = logging.getLogger(__name__) -@dataclass -class _ZPipelineSamplingParams: - """Typed shim satisfying the attributes ZImagePipeline.forward reads - from ``req.sampling_params``. Unlike SimpleNamespace this fails loudly - at construction if a required field is missing.""" - - height: int - width: int - num_inference_steps: int - guidance_scale: float - generator: torch.Generator | None = None - strength: float | None = None - sigmas: list[float] | None = None - max_sequence_length: int = 512 - guidance_rescale: float | None = None - num_outputs_per_prompt: int = 1 - - -@dataclass -class _ZPipelineRequest: - """Typed shim satisfying the attributes ZImagePipeline.forward reads - from ``req``.""" - - request_id: str - sampling_params: _ZPipelineSamplingParams - prompts: list[dict[str, Any]] = field(default_factory=lambda: [{"prompt": "", "negative_prompt": ""}]) - - @property - def num_reqs(self) -> int: - return 1 - - class MingImagePipeline(ZImagePipeline): """Ming-flash-omni-2.0 text-to-image diffusion pipeline. @@ -123,6 +92,10 @@ class MingImagePipeline(ZImagePipeline): self.device = self._execution_device # Ming convention alias self._dtype = dtype + # Request-scoped conditioning handed to the inherited encode_prompt override + self._pending_prompt_embeds: list[torch.Tensor] | None = None + self._pending_negative_prompt_embeds: list[torch.Tensor] | None = None + # Ming's per-checkpoint image-gen configuration. We cannot rely on # ``od_config.hf_config.image_gen_config`` because the diffusion # stage is started with ``hf_config_name: thinker_config`` (the @@ -273,6 +246,21 @@ class MingImagePipeline(ZImagePipeline): latent = self.vae.encode(ref).latent_dist.mode() return (latent - self.vae.config.shift_factor) * self.vae.config.scaling_factor + def encode_prompt(self, *args, **kwargs): # noqa: ARG002 + """Return Ming's precomputed conditioning instead of encoding text. + + NOTE: Ming has no Z-Image text_encoder; its conditioning (cap_feats, optionally ByT5-augmented) + is computed in forward and stashed on `self._pending_*` immediately before + the `super().forward` call, so we simply hand it back here. + """ + if self._pending_prompt_embeds is None: + raise RuntimeError( + "MingImagePipeline.encode_prompt called without pending " + "conditioning; it must run within MingImagePipeline.forward." + ) + + return self._pending_prompt_embeds, self._pending_negative_prompt_embeds + # ------------------------------------------------------------------ # Forward # ------------------------------------------------------------------ @@ -416,17 +404,27 @@ class MingImagePipeline(ZImagePipeline): negative_prompt_embeds = [negative_cap_feats[i] for i in range(negative_cap_feats.shape[0])] else: negative_prompt_embeds = [self.condition_encoder.zero_negative(e) for e in prompt_embeds] + self._pending_prompt_embeds = prompt_embeds + self._pending_negative_prompt_embeds = negative_prompt_embeds - z_sp = _ZPipelineSamplingParams( + # Build a real single-request DiffusionRequestBatch; + # the prompt text itself is a neutral placeholder here. + z_sp = OmniDiffusionSamplingParams( height=height, width=width, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, generator=generator, + output_type="pt", ) - z_req = _ZPipelineRequest( - request_id=req.request_id or "ming-imagegen", - sampling_params=z_sp, + z_req = DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest( + prompt={"prompt": ""}, + sampling_params=z_sp, + request_id=req.request_id or "ming-imagegen", + ) + ] ) # Reference image (img2img) → VAE-encoded latent published on the @@ -446,21 +444,12 @@ class MingImagePipeline(ZImagePipeline): None if ref_latent is None else tuple(ref_latent.shape), ) try: - outputs: DiffusionOutput = super().forward( - z_req, - prompt=None, - height=height, - width=width, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - generator=generator, - output_type="pt", - return_dict=True, - ) + outputs: DiffusionOutput = super().forward(z_req) finally: set_forward_context_ref_latent(None) + # Drop request-scoped conditioning so we don't retain GPU tensors. + self._pending_prompt_embeds = None + self._pending_negative_prompt_embeds = None raw = outputs.output if not isinstance(raw, torch.Tensor): From 51c14eeaede4fb428f456dc785467ce995356518 Mon Sep 17 00:00:00 2001 From: CAICAII <3360776475@qq.com> Date: Sat, 4 Jul 2026 12:16:21 +0800 Subject: [PATCH 4/4] [BugFix] Preserve explicit deploy stage engine_extras (#4255) Signed-off-by: CAICAIIs <3360776475@qq.com> --- tests/helpers/stage_config.py | 13 +++++++ tests/test_config_factory.py | 59 +++++++++++++++++++++++++++++++- vllm_omni/config/stage_config.py | 3 +- 3 files changed, 73 insertions(+), 2 deletions(-) diff --git a/tests/helpers/stage_config.py b/tests/helpers/stage_config.py index 4351eb27d..8706bb1bc 100644 --- a/tests/helpers/stage_config.py +++ b/tests/helpers/stage_config.py @@ -633,6 +633,18 @@ def get_deploy_config_path(rel_path: str) -> str: return str(_DEPLOY_DIR / rel_path) +def get_deploy_config_stage(rel_path: str, stage_id: int) -> dict[str, Any]: + """Return one stage entry from a deploy yaml by ``stage_id``.""" + with open(get_deploy_config_path(rel_path), encoding="utf-8") as f: + cfg = yaml.safe_load(f) or {} + + stage_key = "stages" if "stages" in cfg else "stage_args" + for stage in cfg.get(stage_key, []): + if stage.get("stage_id") == stage_id: + return stage + raise KeyError(f"No stage_id={stage_id} in deploy config {rel_path!r}") + + def _get_config_value_by_path(config_dict: dict, path: str) -> Any: """Read a dot-separated path from a nested dict (e.g. ``engine_args.load_format``).""" current: Any = config_dict @@ -707,6 +719,7 @@ def stage_config_path_for_run_level(stage_config_path: str | None, run_level: st __all__ = [ "get_deploy_config_path", + "get_deploy_config_stage", "modify_stage_config", "stage_config_path_for_run_level", ] diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index 23d36c338..4b956dfcd 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -14,7 +14,7 @@ from unittest.mock import patch import pytest from transformers import PretrainedConfig, Qwen3OmniMoeConfig -from tests.helpers.stage_config import get_deploy_config_path +from tests.helpers.stage_config import get_deploy_config_path, get_deploy_config_stage from vllm_omni.config.config_factory import StageConfigFactory from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline from vllm_omni.config.stage_config import ( @@ -739,6 +739,15 @@ stages: assert deploy.stages[0].compilation_config == {"pass_config": {"fuse_allreduce_rms": False}} assert "compilation_config" not in deploy.stages[0].engine_extras + def test_load_voxcpm2_deploy_config_preserves_engine_extras(self): + deploy_path = get_deploy_config_path("voxcpm2.yaml") + raw_stage = get_deploy_config_stage("voxcpm2.yaml", 0) + expected_runtime_config = raw_stage["engine_extras"]["hf_overrides"]["voxcpm2_runtime_config"] + + deploy = load_deploy_config(deploy_path) + runtime_config = deploy.stages[0].engine_extras["hf_overrides"]["voxcpm2_runtime_config"] + assert runtime_config == expected_runtime_config + def test_merge_pipeline_deploy(self): deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not deploy_path.exists(): @@ -928,6 +937,54 @@ stages: assert "foo" in stage.engine_extras assert stage.engine_extras["foo"] == {"a": 111, "b": {"d": 9, "e": 199}, "c": 3} + def test_explicit_engine_extras_merge_with_flat_and_engine_args(self): + """Explicit engine_extras should be preserved with existing pass-through styles.""" + fake_config = { + "stages": [ + { + "stage_id": 0, + "engine_extras": { + "hf_overrides": { + "runtime_config": { + "explicit_only": True, + "shared": "explicit", + "nested": {"a": 1}, + } + } + }, + "hf_overrides": { + "runtime_config": { + "flat_only": True, + "shared": "flat", + "nested": {"b": 2}, + } + }, + "engine_args": { + "hf_overrides": { + "runtime_config": { + "engine_args_only": True, + "shared": "engine_args", + "nested": {"c": 3}, + } + } + }, + }, + ] + } + + with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config): + deploy = load_deploy_config("dummy.yaml") + + assert deploy.stages[0].engine_extras["hf_overrides"] == { + "runtime_config": { + "engine_args_only": True, + "explicit_only": True, + "flat_only": True, + "nested": {"a": 1, "b": 2, "c": 3}, + "shared": "engine_args", + } + } + def test_deep_merge_does_not_mutate_inputs(self): """Merging engine_args must not mutate the base stage dict.""" fake_config = { diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index f66023bc4..858e9fa15 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -468,6 +468,7 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig: """Parse a single stage entry from deploy YAML into StageDeployConfig.""" # Get the non-reserved keys for this stage flat_args = {k: v for k, v in stage_data.items() if k not in _STAGE_RESERVED_KEYS} + explicit_engine_extras = dict(stage_data.get("engine_extras") or {}) runtime_cfg = dict(stage_data.get("runtime", {})) devices = runtime_cfg.get("devices", stage_data.get("devices")) num_replicas = runtime_cfg.get("num_replicas", stage_data.get("num_replicas", 1)) @@ -495,7 +496,7 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig: kwargs["output_connectors"] = stage_data.get("output_connectors") kwargs["input_connectors"] = stage_data.get("input_connectors") kwargs["default_sampling_params"] = stage_data.get("default_sampling_params") - kwargs["engine_extras"] = flat_args + kwargs["engine_extras"] = _get_recursively_merged_dict(explicit_engine_extras, flat_args) return StageDeployConfig(**kwargs)