From a5574a2fb04ba72303285b94f4ab3ed57682bca2 Mon Sep 17 00:00:00 2001
From: Alex Brooks <albrooks@redhat.com>
Date: Sat, 21 Mar 2026 23:05:18 -0600
Subject: [PATCH] [CI] Add Flux2 Klein Tests (#2027)

Signed-off-by: Alex Brooks <albrooks@redhat.com>
---
 .buildkite/test-merge.yml                     | 24 +++++++
 .../online_serving/test_flux2_expansion.py    | 71 +++++++++++++++++++
 2 files changed, 95 insertions(+)
 create mode 100644 tests/e2e/online_serving/test_flux2_expansion.py

diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml
index b606648c7..2eb68941a 100644
--- a/.buildkite/test-merge.yml
+++ b/.buildkite/test-merge.yml
@@ -126,6 +126,30 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
+  - label: "Diffusion Model Test with L4"
+    timeout_in_minutes: 60
+    depends_on: upload-merge-pipeline
+    # This should run nightly, but moving to the merge pipeline for testing
+    #depends_on: upload-nightly-pipeline
+    # if: build.env("NIGHTLY") == "1"
+    #
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and L4" --run-level "advanced_model"
+    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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
   - label: "Diffusion Tensor Parallelism Test"
     timeout_in_minutes: 20
     depends_on: upload-merge-pipeline
diff --git a/tests/e2e/online_serving/test_flux2_expansion.py b/tests/e2e/online_serving/test_flux2_expansion.py
new file mode 100644
index 000000000..0e9e8c89a
--- /dev/null
+++ b/tests/e2e/online_serving/test_flux2_expansion.py
@@ -0,0 +1,71 @@
+"""
+Tests for Flux2 Klein; currently Dev is implemented separately,
+but ideally these models will fold together in the future.
+"""
+
+import pytest
+
+from tests.conftest import (
+    OmniServer,
+    OmniServerParams,
+    OpenAIClientHandler,
+)
+from tests.utils import hardware_marks
+
+FOUR_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=4)
+POSITIVE_PROMPT = "A cat sitting on a windowsill"
+NEGATIVE_PROMPT = "blurry, low quality"
+
+
+# For now, we only test a good configuration for this model to keep the CI
+# light, but in the near future may explore using random models to check
+# cross-feature compatibility more generally.
+def _get_diffusion_feature_cases(model: str):
+    return [
+        # FP8 / Hybrid sequence parallelism
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--cache-backend",
+                    "cache_dit",
+                    "--ulysses-degree",
+                    "2",
+                    "--quantization",
+                    "fp8",
+                    # NOTE: TP added for test coverage here since it doesn't
+                    # slow things down too much, but fp8 + ulysses SP +
+                    # cache_dit will generally get you the highest speedup here
+                    "--tensor-parallel-size",
+                    "2",
+                ],
+            ),
+            marks=FOUR_CARD_FEATURE_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_diffusion_feature_cases(
+        model="black-forest-labs/FLUX.2-klein-4B",
+    ),
+    indirect=True,
+)
+def test_flux2_klein(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+    request_config = {
+        "model": omni_server.model,
+        "messages": [{"role": "user", "content": POSITIVE_PROMPT}],
+        "extra_body": {
+            "height": 512,
+            "width": 512,
+            "num_inference_steps": 2,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "true_cfg_scale": 4.0,
+            "seed": 42,
+        },
+    }
+
+    openai_client.send_diffusion_request(request_config)