[CI][ROCm] Add Cosmos3 nightly coverage (#7933)

Signed-off-by: haic0 <149741444+haic0@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

.buildkite/amd/scripts/prewarm_cosmos3_aiter.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Prebuild the BF16 AITER MHA modules used by Cosmos3 on ROCm."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import torch
+
+from vllm_omni.diffusion.attention.backends.utils import fa
+from vllm_omni.platforms import current_omni_platform
+
+ATTENTION_VARIANTS = (
+    (True, "mha_fwd_bf16_nbias_mask_nlse_ndropout_nqscale"),
+    (False, "mha_fwd_bf16_nbias_nmask_nlse_ndropout_nqscale"),
+)
+
+
+def main() -> None:
+    if not current_omni_platform.is_rocm():
+        raise RuntimeError("The Cosmos3 AITER prewarm must run on ROCm")
+    if fa.flash_attn_func is None:
+        raise RuntimeError("ROCm FlashAttention/AITER is unavailable")
+
+    cache_dir_value = os.environ.get("AITER_JIT_DIR")
+    if not cache_dir_value:
+        raise RuntimeError("AITER_JIT_DIR must identify the job-local cache")
+    cache_dir = Path(cache_dir_value)
+
+    device = torch.device(current_omni_platform.device_type)
+    dtype = torch.bfloat16
+    batch_size = 1
+    sequence_length = 64
+    num_heads = 8
+    head_dim = 64
+
+    torch.manual_seed(42)
+    query = torch.randn(
+        batch_size,
+        sequence_length,
+        num_heads,
+        head_dim,
+        device=device,
+        dtype=dtype,
+    )
+    key = torch.randn_like(query)
+    value = torch.randn_like(query)
+
+    for causal, expected_module in ATTENTION_VARIANTS:
+        output = fa.flash_attn_func(
+            query,
+            key,
+            value,
+            dropout_p=0.0,
+            softmax_scale=head_dim**-0.5,
+            causal=causal,
+        )
+        if isinstance(output, tuple):
+            output = output[0]
+        current_omni_platform.synchronize()
+
+        if output.shape != query.shape:
+            raise RuntimeError(f"Unexpected AITER output shape: {output.shape}")
+        if not torch.isfinite(output).all():
+            raise RuntimeError("AITER prewarm produced non-finite output")
+
+        module_path = cache_dir / f"{expected_module}.so"
+        if not module_path.is_file():
+            raise RuntimeError(f"Expected AITER module was not built: {module_path}")
+        print(f"Prebuilt Cosmos3 AITER module: {module_path}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

.buildkite/amd/test-amd-nightly.yml

@@ -99,3 +99,36 @@ steps:
             artifact_dir="$${BUILDKITE_BUILD_CHECKOUT_PATH:?}/artifacts/rocm-lingbot-nightly"
             mkdir -p "$$artifact_dir"
             cp -v /tmp/lingbot-pytest.xml "$$artifact_dir/"
+
+  - group: ":full_moon: Cosmos3 MI300 Nightly"
+    steps:
+      - label: "Cosmos3 T2I Function"
+        agent_pool: mi300_1
+        mirror_hardwares: [amdproduction]
+        grade: NonBlocking
+        timeout_in_minutes: 90
+        artifact_paths:
+          - artifacts/rocm-cosmos3-nightly/**/*
+        commands:
+          # Cosmos3 uses both masked and unmasked BF16 AITER MHA modules. Build
+          # them before server startup, then reuse the same job-local cache
+          # during the smoke test.
+          - export VLLM_ROCM_USE_AITER=1
+          - export AITER_JIT_DIR="/tmp/vllm-omni-cosmos3-aiter-$$BUILDKITE_JOB_ID"
+          - export TORCH_EXTENSIONS_DIR="/tmp/vllm-omni-cosmos3-torch-extensions-$$BUILDKITE_JOB_ID"
+          - mkdir -p "$$AITER_JIT_DIR" "$$TORCH_EXTENSIONS_DIR"
+          - pytest --collect-only -q tests/e2e/online_serving/test_cosmos3.py::test_text_to_image_001 -m "core_model and rocm and MI325 and cards_1" --run-level full_model
+          - >-
+            timeout --signal=TERM --kill-after=2m 55m
+            python .buildkite/amd/scripts/prewarm_cosmos3_aiter.py
+          - >-
+            timeout --signal=TERM --kill-after=2m 25m
+            pytest -s -v
+            tests/e2e/online_serving/test_cosmos3.py::test_text_to_image_001
+            -m "core_model and rocm and MI325 and cards_1"
+            --run-level full_model
+            --junitxml=/tmp/cosmos3-pytest.xml
+          - |
+            artifact_dir="$${BUILDKITE_BUILD_CHECKOUT_PATH:?}/artifacts/rocm-cosmos3-nightly"
+            mkdir -p "$$artifact_dir"
+            cp -v /tmp/cosmos3-pytest.xml "$$artifact_dir/"

tests/e2e/online_serving/test_cosmos3.py

@@ -40,7 +40,7 @@ SERVER_ARGS = [
     "Cosmos3OmniDiffusersPipeline",
     "--no-guardrails",
 ]
-SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
+SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100", "rocm": "MI325"})
 
 
 def _get_diffusion_feature_cases(model: str):