diff --git a/.buildkite/amd/scripts/prewarm_cosmos3_aiter.py b/.buildkite/amd/scripts/prewarm_cosmos3_aiter.py new file mode 100644 index 000000000..b7f45fd2a --- /dev/null +++ b/.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() diff --git a/.buildkite/amd/test-amd-nightly.yml b/.buildkite/amd/test-amd-nightly.yml index 9e2b5bacc..1d6775434 100644 --- a/.buildkite/amd/test-amd-nightly.yml +++ b/.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/" diff --git a/tests/e2e/online_serving/test_cosmos3.py b/tests/e2e/online_serving/test_cosmos3.py index f28e2ae26..17fc63077 100644 --- a/tests/e2e/online_serving/test_cosmos3.py +++ b/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):