Merge branch 'main' into governance

.buildkite/test-merge.yml

@@ -119,23 +119,23 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  # - label: "Diffusion Tensor Parallelism Test"
-  #   timeout_in_minutes: 20
-  #   depends_on: upload-merge-pipeline
-  #   commands:
-  #     - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
-  #   agents:
-  #     queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-  #   plugins:
-  #     - docker#v5.2.0:
-  #         image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-  #         always-pull: true
-  #         propagate-environment: true
-  #         shm-size: "8gb"
-  #         environment:
-  #           - "HF_HOME=/fsx/hf_cache"
-  #         volumes:
-  #           - "/fsx/hf_cache:/fsx/hf_cache"
+  - label: "Diffusion Tensor Parallelism Test"
+    timeout_in_minutes: 20
+    depends_on: upload-merge-pipeline
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
+    agents:
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion GPU Worker Test"
     timeout_in_minutes: 20

benchmarks/diffusion/README.md

@@ -43,6 +43,27 @@ The benchmark supports three dataset modes via `--dataset`:
 
 ### VBench dataset
 
+`vbench` only provides prompt data (and image paths for i2v/i2i); it does not carry
+per-request generation fields. In this mode, all requests share CLI values:
+`--width --height --num-frames --fps --num-inference-steps`
+(pass `--width` and `--height` together).
+
+Example (`t2v`):
+
+```bash
+python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+	--base-url http://localhost:8099 \
+	--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+	--task t2v \
+	--dataset vbench \
+	--num-prompts 50 \
+	--width 640 --height 480 \
+	--num-frames 81 --fps 16 \
+	--num-inference-steps 40
+```
+
+Note: `vbench` can also be used for other tasks such as `t2i` / `i2v` (and `i2i`). For `t2i`, the loader reuses VBench t2v text prompts; for `i2v` / `i2i`, it loads the VBench i2v dataset (with image paths).
+
 If you use i2v/i2i bench datasets and need auto-download support, you may need:
 
 ```bash
@@ -90,7 +111,7 @@ Common optional flags:
 Related flags: `--width`, `--height`, `--num-frames`, `--fps`, `--num-inference-steps`.
 
 - For `vbench` / `random`: these CLI flags act as global defaults for all generated requests.
-- For `trace`: each request can carry its own fields (e.g., `width/height/num_frames/num_inference_steps`).
+- For `trace`: requests can carry their own fields (e.g., `width/height/num_frames/num_inference_steps`), with overrides/fallbacks as below.
 
 Precedence rules for `trace` (i.e., what actually gets sent):
 

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -256,8 +256,14 @@ class DiffusersPipelineLoader:
         load_device: str,
         load_format: str = "default",
         custom_pipeline_name: str | None = None,
+        device: torch.device | None = None,
     ) -> nn.Module:
         """Load a model with the given configurations."""
+        # CPU offload + FP8: load weights on device for FP8 quantization
+        if load_device == "cpu" and od_config.quantization and od_config.quantization.lower() != "none":
+            load_device = device.type
+            logger.info(f"Quantization enabled with CPU offload, using {load_device} for weight loading")
+
         target_device = torch.device(load_device)
         with set_default_torch_dtype(od_config.dtype):
             if od_config.parallel_config.use_hsdp:

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -127,6 +127,7 @@ class DiffusionModelRunner:
                     load_device=load_device,
                     load_format=load_format,
                     custom_pipeline_name=custom_pipeline_name,
+                    device=self.device,
                 )
         time_after_load = time.perf_counter()
 

vllm_omni/entrypoints/async_omni.py

@@ -225,6 +225,7 @@ class AsyncOmni(OmniBase):
                     "enforce_eager": kwargs.get("enforce_eager", False),
                     "diffusion_load_format": kwargs.get("diffusion_load_format", "default"),
                     "custom_pipeline_args": kwargs.get("custom_pipeline_args", None),
+                    "quantization": kwargs.get("quantization", None),
                     "worker_extension_cls": kwargs.get("worker_extension_cls", None),
                     "enable_sleep_mode": kwargs.get("enable_sleep_mode", False),
                     "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True),