diff --git a/.buildkite/amd/test-amd-ready.yml b/.buildkite/amd/test-amd-ready.yml index 3860df345..245eb2cea 100644 --- a/.buildkite/amd/test-amd-ready.yml +++ b/.buildkite/amd/test-amd-ready.yml @@ -128,6 +128,15 @@ steps: tests/diffusion/attention/ -m 'core_model and rocm and MI325 and cards_1' --run-level "core_model" + # MammothModa2 exercises the same two AITER MHA modules built by the + # preceding suite. Keep its default-backend run in this job so a cold + # worker compiles each module once instead of timing out the generic + # model job. + - >- + timeout --signal=TERM --kill-after=2m 10m + pytest -s -v -n 1 + tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py + --run-level "core_model" # AR paged attention is real GPU work, including cold AITER compilation. # Keep its complete numerical matrix out of the CPU-marked Simple shards. @@ -216,9 +225,17 @@ steps: - >- timeout 15m pytest -sv tests/diffusion/models/ --ignore=tests/diffusion/models/auk/test_auk_vae_fastpath.py + --ignore=tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py -m 'core_model and cuda and not (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)' --run-level "core_model" + # Keep the model-specific numerical coverage blocking without paying + # the default ROCm backend's cold AITER compilation cost in this lane. + - >- + timeout 10m env DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA + pytest -sv + tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py + --run-level "core_model" - label: "Diffusion ยท Wan22 Test" agent_pool: mi300_1 diff --git a/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py b/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py index cec4fa316..fc7786929 100644 --- a/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py +++ b/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py @@ -3,6 +3,8 @@ """Real-shape checks for MammothModa2's fused QK norm/RoPE and shared CUDA attention backend against the native shared RMSNorm path.""" +import os + import pytest import torch from diffusers.models.attention_processor import Attention @@ -10,6 +12,9 @@ from diffusers.models.attention_processor import Attention import vllm_omni.diffusion.attention.backends.sdpa as sdpa_backend import vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model as mammoth_dit from tests.helpers.mark import hardware_marks +from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum +from vllm_omni.diffusion.config import set_current_diffusion_config +from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.layers.norm import RMSNorm from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import TransformerBlock @@ -26,6 +31,13 @@ pytestmark = [ DIM, HEADS, KV_HEADS = 2520, 21, 7 +@pytest.fixture(autouse=True) +def _set_diffusion_config(): + """Construct attention layers with the same config path used in production.""" + with set_current_diffusion_config(OmniDiffusionConfig()): + yield + + def _norm_attention(head_dim, dtype): attn = Attention(query_dim=head_dim * HEADS, heads=HEADS, kv_heads=KV_HEADS, dim_head=head_dim) attn.norm_q = RMSNorm(head_dim) @@ -33,6 +45,24 @@ def _norm_attention(head_dim, dtype): return attn.to(device="cuda", dtype=dtype) +def test_explicit_attention_backend_environment_reaches_dit_layer(): + backend_name = os.environ.get("DIFFUSION_ATTENTION_BACKEND") + if backend_name is None or backend_name.lower() == "auto": + pytest.skip("no explicit diffusion attention backend") + assert backend_name is not None + + block = TransformerBlock( + 48, + 6, + 2, + multiple_of=8, + ffn_dim_multiplier=1.0, + norm_eps=1e-5, + ) + expected_backend = DiffusionAttentionBackendEnum[backend_name.upper()].get_class() + assert block.attn.omni_attn.attn_backend is expected_backend + + def test_unpaired_rope_preserves_native_arithmetic(monkeypatch): """A generic block accepts independently valued even and odd RoPE lanes.""" torch.manual_seed(42)