Merge remote-tracking branch 'remotes/origin/main' into dev/vllm-align

.buildkite/test-merge.yml

@@ -12,6 +12,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -29,6 +30,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -46,6 +48,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -64,6 +67,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -81,6 +85,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -98,6 +103,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -116,6 +122,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -134,6 +141,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -152,6 +160,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -178,6 +187,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -197,6 +207,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -226,6 +237,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -261,6 +277,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -272,37 +293,50 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  # - label: "Bagel Text2Img Model Test with H100"
-  #   timeout_in_minutes: 30
-  #   depends_on: upload-merge-pipeline
-  #   commands:
-  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #     - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
-  #   agents:
-  #     queue: "mithril-h100-pool"
-  #   plugins:
-  #     - kubernetes:
-  #         podSpec:
-  #           containers:
-  #             - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-  #               resources:
-  #                 limits:
-  #                   nvidia.com/gpu: 1
-  #               volumeMounts:
-  #                 - name: devshm
-  #                   mountPath: /dev/shm
-  #                 - name: hf-cache
-  #                   mountPath: /root/.cache/huggingface
-  #               env:
-  #                 - name: HF_HOME
-  #                   value: /root/.cache/huggingface
-  #           nodeSelector:
-  #             node.kubernetes.io/instance-type: gpu-h100-sxm
-  #           volumes:
-  #             - name: devshm
-  #               emptyDir:
-  #                 medium: Memory
-  #             - name: hf-cache
-  #               hostPath:
-  #                 path: /mnt/hf-cache
-  #                 type: DirectoryOrCreate
+  - label: "Bagel Model Test with H100 (Real Weights)"
+    timeout_in_minutes: 60
+    depends_on: upload-merge-pipeline
+    commands:
+      - |
+        timeout 55m bash -c '
+          set -e
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_TEST_CLEAN_GPU_MEMORY=1
+          export VLLM_IMAGE_FETCH_TIMEOUT=60
+          pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory"
+          pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model"
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate

.buildkite/test-nightly.yml

@@ -30,6 +30,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -58,6 +63,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -86,6 +92,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -116,6 +127,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -149,6 +161,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:

.buildkite/test-ready.yml

@@ -12,6 +12,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -28,6 +29,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -44,6 +46,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -60,6 +63,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -76,6 +80,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -93,6 +98,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -110,6 +116,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -137,6 +144,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -159,6 +167,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -221,6 +230,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -268,7 +278,7 @@ steps:
         timeout 30m bash -c '
           export VLLM_WORKER_MULTIPROC_METHOD=spawn
           export VLLM_TEST_CLEAN_GPU_MEMORY=1
-          pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
+          pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model"
         '
     agents:
       queue: "mithril-h100-pool"
@@ -288,6 +298,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -306,7 +321,7 @@ steps:
         timeout 30m bash -c '
           export VLLM_WORKER_MULTIPROC_METHOD=spawn
           export VLLM_TEST_CLEAN_GPU_MEMORY=1
-          pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py
+          pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model"
         '
     agents:
       queue: "mithril-h100-pool"
@@ -326,6 +341,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -345,7 +365,7 @@ steps:
           export VLLM_WORKER_MULTIPROC_METHOD=spawn
           export VLLM_TEST_CLEAN_GPU_MEMORY=1
           export VLLM_IMAGE_FETCH_TIMEOUT=60
-          pytest -s -v tests/e2e/online_serving/test_bagel_online.py
+          pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model"
         '
     agents:
       queue: "mithril-h100-pool"
@@ -365,6 +385,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:

tests/conftest.py

@@ -1236,15 +1236,14 @@ def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: st
         port = params.port
         stage_config_path = params.stage_config_path
         if run_level == "advanced_model" and stage_config_path is not None:
+            # Dynamically detect stages from config to avoid KeyError
+            # for models with fewer stages (e.g., BAGEL has 2, Qwen3-Omni has 3)
+            with open(stage_config_path, encoding="utf-8") as f:
+                _cfg = yaml.safe_load(f) or {}
+            _stage_ids = [s["stage_id"] for s in _cfg.get("stage_args", []) if "stage_id" in s]
             stage_config_path = modify_stage_config(
                 stage_config_path,
-                deletes={
-                    "stage_args": {
-                        0: ["engine_args.load_format"],
-                        1: ["engine_args.load_format"],
-                        2: ["engine_args.load_format"],
-                    }
-                },
+                deletes={"stage_args": {sid: ["engine_args.load_format"] for sid in _stage_ids}},
             )
 
         server_args = params.server_args or []

tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml

@@ -21,6 +21,7 @@ stage_args:
       enable_prefix_caching: false
       max_num_batched_tokens: 32768
       tensor_parallel_size: 1
+      load_format: dummy
       omni_kv_config:
         need_send_cache: true
         kv_transfer_criteria:
@@ -54,6 +55,7 @@ stage_args:
       enable_prefix_caching: false
       max_num_batched_tokens: 32768
       tensor_parallel_size: 1
+      load_format: dummy
       omni_kv_config:
         need_recv_cache: true
     engine_input_source: [0]

tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml

@@ -21,6 +21,7 @@ stage_args:
       enable_prefix_caching: false
       max_num_batched_tokens: 32768
       tensor_parallel_size: 1
+      load_format: dummy
       omni_kv_config:
         need_send_cache: true
         kv_transfer_criteria:
@@ -53,6 +54,7 @@ stage_args:
       enable_prefix_caching: false
       max_num_batched_tokens: 32768
       tensor_parallel_size: 1
+      load_format: dummy
       omni_kv_config:
         need_recv_cache: true
     engine_input_source: [0]

tests/e2e/offline_inference/test_bagel_img2img.py

@@ -22,6 +22,7 @@ import pytest
 from PIL import Image
 from vllm.assets.image import ImageAsset
 
+from tests.conftest import modify_stage_config
 from tests.utils import hardware_test
 from vllm_omni.entrypoints.omni import Omni
 
@@ -168,17 +169,39 @@ def _generate_bagel_img2img(
     return generated_image
 
 
+def _resolve_stage_config(config_path: str, run_level: str) -> str:
+    """Resolve stage config based on run level.
+
+    For advanced_model (real weights), strip load_format: dummy so the model
+    falls back to loading real weights from HuggingFace.
+    """
+    if run_level == "advanced_model":
+        return modify_stage_config(
+            config_path,
+            deletes={
+                "stage_args": {
+                    0: ["engine_args.load_format"],
+                    1: ["engine_args.load_format"],
+                }
+            },
+        )
+    return config_path
+
+
 @pytest.mark.core_model
+@pytest.mark.advanced_model
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "H100"})
-def test_bagel_img2img_shared_memory_connector():
+def test_bagel_img2img_shared_memory_connector(run_level):
     """Test Bagel img2img with shared memory connector."""
     input_image = _load_input_image()
     config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
+    config_path = _resolve_stage_config(config_path, run_level)
     omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path, stage_init_timeout=300)
 
     try:
         generated_image = _generate_bagel_img2img(omni, input_image)
-        _validate_pixels(generated_image)
+        if run_level == "advanced_model":
+            _validate_pixels(generated_image)
     finally:
         omni.close()

tests/e2e/offline_inference/test_bagel_text2img.py

@@ -28,6 +28,7 @@ from typing import Any
 import pytest
 from PIL import Image
 
+from tests.conftest import modify_stage_config
 from tests.utils import hardware_test
 from vllm_omni.entrypoints.omni import Omni
 
@@ -158,17 +159,39 @@ def _generate_bagel_image(omni: Omni, prompt: str = DEFAULT_PROMPT) -> Image.Ima
     return generated_image
 
 
+def _resolve_stage_config(config_path: str, run_level: str) -> str:
+    """Resolve stage config based on run level.
+
+    For advanced_model (real weights), strip load_format: dummy so the model
+    falls back to loading real weights from HuggingFace.
+    """
+    if run_level == "advanced_model":
+        return modify_stage_config(
+            config_path,
+            deletes={
+                "stage_args": {
+                    0: ["engine_args.load_format"],
+                    1: ["engine_args.load_format"],
+                }
+            },
+        )
+    return config_path
+
+
 @pytest.mark.core_model
+@pytest.mark.advanced_model
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "H100"})
-def test_bagel_text2img_shared_memory_connector():
+def test_bagel_text2img_shared_memory_connector(run_level):
     """Test Bagel text2img with shared memory connector."""
     config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
+    config_path = _resolve_stage_config(config_path, run_level)
     omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path, stage_init_timeout=300)
 
     try:
         generated_image = _generate_bagel_image(omni)
-        _validate_pixels(generated_image)
+        if run_level == "advanced_model":
+            _validate_pixels(generated_image)
     finally:
         omni.close()
 
@@ -251,9 +274,10 @@ def _load_mooncake_config(host: str, rpc_port: int, http_port: int) -> str:
 
 
 @pytest.mark.core_model
+@pytest.mark.advanced_model
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "H100"})
-def test_bagel_text2img_mooncake_connector():
+def test_bagel_text2img_mooncake_connector(run_level):
     """Test Bagel text2img with Mooncake connector for inter-stage communication."""
     MOONCAKE_HOST = "127.0.0.1"
     MOONCAKE_RPC_PORT = _find_free_port()
@@ -291,10 +315,12 @@ def test_bagel_text2img_mooncake_connector():
             http_port=MOONCAKE_HTTP_PORT,
         )
 
+        temp_config_file = _resolve_stage_config(temp_config_file, run_level)
         omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=temp_config_file, stage_init_timeout=300)
 
         generated_image = _generate_bagel_image(omni)
-        _validate_pixels(generated_image)
+        if run_level == "advanced_model":
+            _validate_pixels(generated_image)
 
     finally:
         if omni:

tests/e2e/online_serving/test_bagel_online.py

@@ -22,22 +22,18 @@ Equivalent to running:
 
 import base64
 import os
-import signal
-import socket
-import subprocess
-import sys
-import time
 from io import BytesIO
 from pathlib import Path
-from typing import Any
 
 import pytest
-import requests
-from PIL import Image
 from vllm.assets.image import ImageAsset
 
+from tests.conftest import OmniServerParams
 from tests.utils import hardware_test
 
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
 MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
 STAGE_CONFIGS_PATH = str(
     Path(__file__).parent.parent / "offline_inference" / "stage_configs" / "bagel_sharedmemory_ci.yaml"
@@ -46,202 +42,74 @@ STAGE_CONFIGS_PATH = str(
 TEXT2IMG_PROMPT = "A cute cat"
 IMG2IMG_PROMPT = "Change the grass color to red"
 
-
-class BagelOmniServer:
-    """Context manager to start/stop a vLLM-Omni server for Bagel model tests."""
-
-    def __init__(
-        self,
-        model: str = MODEL,
-        stage_configs_path: str = STAGE_CONFIGS_PATH,
-        env_dict: dict[str, str] | None = None,
-    ) -> None:
-        self.model = model
-        self.stage_configs_path = stage_configs_path
-        self.env_dict = env_dict
-        self.proc: subprocess.Popen | None = None
-        self.host = "127.0.0.1"
-        self.port = _find_free_port()
-
-    @property
-    def base_url(self) -> str:
-        return f"http://{self.host}:{self.port}"
-
-    def _start_server(self) -> None:
-        env = os.environ.copy()
-        if self.env_dict is not None:
-            env.update(self.env_dict)
-
-        cmd = [
-            sys.executable,
-            "-m",
-            "vllm_omni.entrypoints.cli.main",
-            "serve",
-            self.model,
-            "--omni",
-            "--host",
-            self.host,
-            "--port",
-            str(self.port),
-            "--stage-configs-path",
-            self.stage_configs_path,
-            "--stage-init-timeout",
-            "300",
-        ]
-
-        self.proc = subprocess.Popen(
-            cmd,
-            env=env,
-            start_new_session=True,
-        )
-
-        try:
-            if not _wait_for_port(self.host, self.port, timeout=600, proc=self.proc):
-                self.terminate()
-                raise RuntimeError(f"Server failed to start within 600 seconds on {self.host}:{self.port}")
-        except Exception:
-            self.terminate()
-            raise
-
-    def __enter__(self):
-        self._start_server()
-        return self
-
-    def terminate(self) -> None:
-        if self.proc:
-            try:
-                os.killpg(os.getpgid(self.proc.pid), signal.SIGTERM)
-            except ProcessLookupError:
-                pass
-            try:
-                self.proc.wait(timeout=30)
-            except subprocess.TimeoutExpired:
-                try:
-                    os.killpg(os.getpgid(self.proc.pid), signal.SIGKILL)
-                except ProcessLookupError:
-                    pass
-                self.proc.wait()
-            self.proc = None
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.terminate()
+# Create parameter combinations for model and stage config
+test_params = [
+    OmniServerParams(
+        model=MODEL,
+        stage_config_path=STAGE_CONFIGS_PATH,
+        server_args=["--stage-init-timeout", "300"],
+    ),
+]
 
 
-def _find_free_port() -> int:
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(("127.0.0.1", 0))
-        s.listen(1)
-        return s.getsockname()[1]
+def _build_text2img_messages(prompt: str) -> list[dict]:
+    """Build OpenAI-format messages for text2img generation."""
+    return [
+        {
+            "role": "user",
+            "content": [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}],
+        }
+    ]
 
 
-def _wait_for_port(host: str, port: int, timeout: int = 600, proc: subprocess.Popen | None = None) -> bool:
-    start = time.time()
-    while time.time() - start < timeout:
-        if proc is not None and proc.poll() is not None:
-            # Server process exited early
-            return False
-        try:
-            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
-                sock.settimeout(1)
-                if sock.connect_ex((host, port)) == 0:
-                    return True
-        except Exception:
-            pass
-        time.sleep(2)
-    return False
+def _build_img2img_messages(prompt: str, image_b64: str) -> list[dict]:
+    """Build OpenAI-format messages for img2img generation."""
+    return [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"},
+                {
+                    "type": "image_url",
+                    "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"},
+                },
+            ],
+        }
+    ]
 
 
-def _send_chat_request(
-    server_url: str,
-    prompt: str,
-    *,
-    modality: str = "text2img",
-    image: Image.Image | None = None,
-    timeout: int = 300,
-) -> dict[str, Any]:
-    """Send a chat completion request matching the openai_chat_client.py format."""
-    content: list[dict[str, Any]] = [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}]
-
-    if image is not None:
-        buffer = BytesIO()
-        image.save(buffer, format="JPEG")
-        b64_data = base64.b64encode(buffer.getvalue()).decode("utf-8")
-        content.append(
-            {
-                "type": "image_url",
-                "image_url": {"url": f"data:image/jpeg;base64,{b64_data}"},
-            }
-        )
-
-    payload: dict[str, Any] = {
-        "messages": [{"role": "user", "content": content}],
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_bagel_text2img_online(omni_server, openai_client) -> None:
+    """Test Bagel text2img via OpenAI-compatible chat completions API."""
+    request_config = {
+        "model": omni_server.model,
+        "messages": _build_text2img_messages(TEXT2IMG_PROMPT),
+        "modalities": ["image"],
     }
 
-    if modality in ("text2img", "img2img"):
-        payload["modalities"] = ["image"]
-
-    resp = requests.post(
-        f"{server_url}/v1/chat/completions",
-        headers={"Content-Type": "application/json"},
-        json=payload,
-        timeout=timeout,
-    )
-    resp.raise_for_status()
-    return resp.json()
-
-
-def _extract_image_from_response(data: dict[str, Any]) -> Image.Image | None:
-    """Extract the generated PIL Image from a chat completion response."""
-    for choice in data.get("choices", []):
-        content = choice.get("message", {}).get("content")
-        if isinstance(content, list) and content:
-            first_item = content[0]
-            if isinstance(first_item, dict) and "image_url" in first_item:
-                url = first_item["image_url"].get("url", "")
-                if url.startswith("data:image"):
-                    _, b64 = url.split(",", 1)
-                    return Image.open(BytesIO(base64.b64decode(b64)))
-    return None
+    openai_client.send_diffusion_request(request_config)
 
 
 @pytest.mark.core_model
+@pytest.mark.advanced_model
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "H100"})
-def test_bagel_text2img_online():
-    """Test Bagel text2img via OpenAI-compatible chat completions API."""
-    with BagelOmniServer() as server:
-        response_data = _send_chat_request(
-            server.base_url,
-            TEXT2IMG_PROMPT,
-            modality="text2img",
-        )
-
-        image = _extract_image_from_response(response_data)
-        assert image is not None, f"No image in response: {response_data}"
-        image.load()
-
-        w, h = image.size
-        assert w > 0 and h > 0, f"Invalid image size: {image.size}"
-
-
-@pytest.mark.core_model
-@pytest.mark.diffusion
-@hardware_test(res={"cuda": "H100"})
-def test_bagel_img2img_online():
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_bagel_img2img_online(omni_server, openai_client) -> None:
     """Test Bagel img2img via OpenAI-compatible chat completions API."""
     input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
+    buffer = BytesIO()
+    input_image.save(buffer, format="JPEG")
+    image_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
 
-    with BagelOmniServer() as server:
-        response_data = _send_chat_request(
-            server.base_url,
-            IMG2IMG_PROMPT,
-            modality="img2img",
-            image=input_image,
-        )
+    request_config = {
+        "model": omni_server.model,
+        "messages": _build_img2img_messages(IMG2IMG_PROMPT, image_b64),
+        "modalities": ["image"],
+    }
 
-        image = _extract_image_from_response(response_data)
-        assert image is not None, f"No image in response: {response_data}"
-        image.load()
-
-        w, h = image.size
-        assert w > 0 and h > 0, f"Invalid image size: {image.size}"
+    openai_client.send_diffusion_request(request_config)

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py

@@ -95,6 +95,9 @@ def get_qwen_image_layered_pre_process_func(
             else:
                 image = cast(PIL.Image.Image | torch.Tensor | np.ndarray, raw_image)
 
+            if isinstance(image, PIL.Image.Image) and image.mode != "RGBA":
+                image = image.convert("RGBA")
+
             # 1. calculate dimensions
             image_size = image.size
             assert request.sampling_params.resolution in [640, 1024], (
@@ -652,6 +655,8 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
             width = req.sampling_params.width
         else:
             # fallback to run pre-processing in pipeline (debug only)
+            if isinstance(image, PIL.Image.Image) and image.mode != "RGBA":
+                image = image.convert("RGBA")
             image_size = image[0].size if isinstance(image, list) else image.size
             assert resolution in [640, 1024], f"resolution must be either 640 or 1024, but got {resolution}"
             calculated_width, calculated_height = calculate_dimensions(

vllm_omni/diffusion/request.py

@@ -28,17 +28,23 @@ class OmniDiffusionRequest:
 
     def __post_init__(self):
         """Initialize dependent fields after dataclass initialization."""
+        # Detect whether user explicitly provided guidance_scale.
+        # The sentinel default is 0.0 (false-like); any truthy value means
+        # the caller set it intentionally.  We must resolve this BEFORE
+        # auto-filling guidance_scale_2, otherwise the sentinel leaks into
+        # guidance_scale_2.
+        if self.sampling_params.guidance_scale:
+            self.sampling_params.guidance_scale_provided = True
+        else:
+            self.sampling_params.guidance_scale = 1.0
+
         # Set do_classifier_free_guidance based on guidance scale and negative prompt
         if self.sampling_params.guidance_scale > 1.0 and any(
             (not isinstance(p, str) and p.get("negative_prompt")) for p in self.prompts
         ):
             self.sampling_params.do_classifier_free_guidance = True
+
+        # Auto-fill guidance_scale_2 from the (now-resolved) guidance_scale
+        # so downstream code always has a valid value.
         if self.sampling_params.guidance_scale_2 is None:
             self.sampling_params.guidance_scale_2 = self.sampling_params.guidance_scale
-
-        # The dataclass default value is 0 (false-like), used to detect whether user explicitly provides this value
-        # After this check is done, reset this value to old default 1
-        if self.sampling_params.guidance_scale:
-            self.sampling_params.guidance_scale_provided = True
-        else:
-            self.sampling_params.guidance_scale = 1.0

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -221,6 +221,7 @@ class DiffusionModelRunner:
                 not getattr(req, "skip_cache_refresh", False)
                 and self.cache_backend is not None
                 and self.cache_backend.is_enabled()
+                and req.sampling_params.num_inference_steps is not None
             ):
                 self.cache_backend.refresh(self.pipeline, req.sampling_params.num_inference_steps)
 

vllm_omni/entrypoints/async_omni_diffusion.py

@@ -172,9 +172,6 @@ class AsyncOmniDiffusion:
         if request_id is None:
             request_id = f"diff-{uuid.uuid4().hex[:16]}"
 
-        if sampling_params.guidance_scale:
-            sampling_params.guidance_scale_provided = True
-
         if lora_request is not None:
             sampling_params.lora_request = lora_request
 

vllm_omni/entrypoints/openai/serving_chat.py

@@ -2052,20 +2052,26 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 except ValueError:
                     logger.warning("Invalid size format: %s", extra_body.get("size"))
 
-            # Get request parameters from extra_body
-            # Text-to-image parameters (ref: text_to_image.py)
-            num_inference_steps = extra_body.get("num_inference_steps", 50)
+            # Get request parameters from extra_body.
+            # Avoid hardcoded defaults here — let each pipeline's forward()
+            # method apply its own model-specific default when the user does
+            # not provide a value.
+            num_inference_steps = extra_body.get("num_inference_steps")
             guidance_scale = extra_body.get("guidance_scale")
-            true_cfg_scale = extra_body.get("true_cfg_scale")  # Qwen-Image specific
+            true_cfg_scale = extra_body.get("true_cfg_scale") or extra_body.get("cfg_scale")
             seed = extra_body.get("seed")
             negative_prompt = extra_body.get("negative_prompt")
             num_outputs_per_prompt = extra_body.get("num_outputs_per_prompt", 1)
 
             # Text-to-video parameters (ref: text_to_video.py)
             num_frames = extra_body.get("num_frames")
-            guidance_scale_2 = extra_body.get("guidance_scale_2")  # For video high-noise CFG
+            guidance_scale_2 = extra_body.get("guidance_scale_2")
             lora_body = extra_body.get("lora")
 
+            # Qwen-Image-Layered parameters
+            layers = extra_body.get("layers")
+            resolution = extra_body.get("resolution")
+
             logger.info(
                 "Diffusion chat request %s: prompt=%r, ref_images=%d, params=%s",
                 request_id,
@@ -2089,25 +2095,27 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 "negative_prompt": negative_prompt,
             }
             gen_params = OmniDiffusionSamplingParams(
-                num_inference_steps=num_inference_steps,
                 height=height,
                 width=width,
                 num_outputs_per_prompt=num_outputs_per_prompt,
                 seed=seed,
             )
 
+            # Only override defaults when the user explicitly provides values
+            if num_inference_steps is not None:
+                gen_params.num_inference_steps = num_inference_steps
             if guidance_scale is not None:
                 gen_params.guidance_scale = guidance_scale
-
-            # Add Qwen-Image specific parameter
             if true_cfg_scale is not None:
                 gen_params.true_cfg_scale = true_cfg_scale
-
-            # Add video generation parameters if set
             if num_frames is not None:
                 gen_params.num_frames = num_frames
             if guidance_scale_2 is not None:
                 gen_params.guidance_scale_2 = guidance_scale_2
+            if layers is not None:
+                gen_params.layers = layers
+            if resolution is not None:
+                gen_params.resolution = resolution
 
             # Parse per-request LoRA (works for both AsyncOmniDiffusion and AsyncOmni).
             if lora_body and isinstance(lora_body, dict):

vllm_omni/inputs/data.py

@@ -234,8 +234,9 @@ class OmniDiffusionSamplingParams:
     step_index: int | None = None
     boundary_ratio: float | None = None
 
-    # Scheduler parameters
-    num_inference_steps: int = 50
+    # Scheduler parameters – ``None`` means "not explicitly set by the caller";
+    # each pipeline's ``forward()`` decides its own model-specific default.
+    num_inference_steps: int | None = None
     guidance_scale: float = 0.0
     guidance_scale_provided: bool = False
     guidance_scale_2: float | None = None

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -396,6 +396,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # Keys that should stay on GPU in model_intermediate_buffer to avoid
         # CPU-to-GPU round-trips on every decode step.
         self.gpu_resident_buffer_keys: set[str] = {
+            "audio_codes",
             "last_talker_hidden",
             "tts_pad_embed",
             "tailing_text_hidden",
@@ -650,7 +651,8 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # Stays on GPU - gpu_resident_buffer_keys avoids the CPU round-trip.
         if hidden_states.numel() == 0:
             return {}
-        return {"last_talker_hidden": hidden_states[-1, :].detach()}
+        last = hidden_states[-1, :].detach()
+        return {"last_talker_hidden": last}
 
     # -------------------- prompt construction helpers --------------------
 

vllm_omni/worker/gpu_ar_model_runner.py

@@ -560,6 +560,29 @@ class GPUARModelRunner(OmniGPUModelRunner):
             hidden_states, multimodal_outputs, num_scheduled_tokens_np, scheduler_output
         )
 
+        # Pre-copy multimodal tensors to CPU once (not per-request) to avoid
+        # redundant D2H transfers when gpu_resident_buffer_keys keeps them on GPU.
+        mm_cpu: dict[str, object] = {}
+        if isinstance(multimodal_outputs, dict) and multimodal_outputs:
+            for k, v in multimodal_outputs.items():
+                try:
+                    if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]:
+                        mm_cpu[k] = v.detach().to("cpu").contiguous()
+                    elif isinstance(v, dict):
+                        sub_dict: dict[str, torch.Tensor] = {}
+                        for sk, sv in v.items():
+                            if isinstance(sv, torch.Tensor) and sv.shape[0] == hidden_states_cpu.shape[0]:
+                                sub_dict[str(sk)] = sv.detach().to("cpu").contiguous()
+                        if sub_dict:
+                            mm_cpu[k] = sub_dict
+                    elif isinstance(v, list):
+                        element = v[0]
+                        if isinstance(element, torch.Tensor):
+                            element = element.detach().to("cpu").contiguous()
+                        mm_cpu[k] = element
+                except Exception as e:
+                    logger.error(f"Error in merge multimodal outputs: {e}")
+
         pooler_output: list[dict[str, object]] = []
         for rid in req_ids_output_copy:
             idx = req_id_to_index_output_copy[rid]
@@ -568,28 +591,20 @@ class GPUARModelRunner(OmniGPUModelRunner):
             end = start + sched
             hidden_slice = hidden_states_cpu[start:end]
             payload: dict[str, object] = {"hidden": hidden_slice}
-            if isinstance(multimodal_outputs, dict) and multimodal_outputs:
+            if mm_cpu:
                 mm_payload: dict[str, object] = {}
-                for k, v in multimodal_outputs.items():
-                    try:
-                        if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]:
-                            mm_payload[k] = v.detach().to("cpu")[start:end].contiguous()
-                        elif isinstance(v, dict):
-                            sub_dict: dict[str, torch.Tensor] = {}
-                            for sk, sv in v.items():
-                                if isinstance(sv, torch.Tensor) and sv.shape[0] == hidden_states_cpu.shape[0]:
-                                    sub_dict[str(sk)] = sv.detach().to("cpu")[start:end].contiguous()
-                            if sub_dict:
-                                mm_payload[k] = sub_dict
-                        elif isinstance(v, list):
-                            element = v[0]
-                            if isinstance(element, torch.Tensor):
-                                element = element.detach().to("cpu").contiguous()
-                            mm_payload[k] = element
-                    except Exception as e:
-                        logger.error(f"Error in merge multimodal outputs: {e}")
-                if mm_payload:
-                    payload.update(mm_payload)
+                for k, v in mm_cpu.items():
+                    if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]:
+                        mm_payload[k] = v[start:end].contiguous()
+                    elif isinstance(v, dict):
+                        mm_payload[k] = {sk: sv[start:end].contiguous() for sk, sv in v.items()}
+                    elif isinstance(v, torch.Tensor):
+                        # List-derived tensor payloads are request-invariant; clone to
+                        # avoid accidental cross-request aliasing on downstream mutation.
+                        mm_payload[k] = v.clone()
+                    else:
+                        mm_payload[k] = v
+                payload.update(mm_payload)
             pooler_output.append(payload)
         with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):
             if self.routed_experts_initialized:

vllm_omni/worker/gpu_model_runner.py

@@ -1313,14 +1313,16 @@ class OmniGPUModelRunner(GPUModelRunner):
             None, self.vllm_config, cudagraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc
         ):
             req_embeds, code_predictor_codes = self.talker_mtp(req_input_ids, req_embeds, last_talker_hidden, text_step)
-        # update the inputs_embeds and code_predictor_codes
-        code_predictor_codes_cpu = code_predictor_codes.detach().to("cpu").contiguous()
+        # code_predictor_codes stays on GPU here; _update_intermediate_buffer
+        # keeps it device-resident when the key is in gpu_resident_buffer_keys.
+        # D2H is deferred to sample_tokens where hidden_states.to("cpu") already
+        # syncs the stream, avoiding a per-step cudaStreamSynchronize.
         out_key = getattr(self.model, "talker_mtp_output_key", "code_predictor_codes")
         for idx, req_id in enumerate(decode_req_ids):
             req_index = self.input_batch.req_ids.index(req_id)
             start_offset = int(self.query_start_loc.cpu[req_index])
             inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1]
-            update_dict = {out_key: code_predictor_codes_cpu[idx : idx + 1]}
+            update_dict = {out_key: code_predictor_codes[idx : idx + 1]}
             self._merge_additional_information_update(req_id, update_dict)
 
     def _model_forward(