[CI/Build][ROCm] Reuse AITER cache for Mammoth attention tests (#8527)

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

.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

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)