sync: rebase PR 8165 routing fix onto current main

.buildkite/amd/test-amd-merge.yml

@@ -41,7 +41,9 @@ steps:
         - export VLLM_ROCM_USE_AITER=0
         # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
         # that is not supported on CPU
-        - "pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
+        # Keep multi-GPU tests out of this one-GPU pool even when their module
+        # also carries the cpu marker for unit tests.
+        - "pytest -sv tests/diffusion -m 'core_model and cpu and not (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)' --ignore=tests/diffusion/cache/test_teacache_extractors.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
 
     - label: "Simple · Engine&Entrypoints Test"
       agent_pool: mi300_1
@@ -131,6 +133,11 @@ steps:
       grade: Blocking
       commands:
         - pytest -s -v tests/diffusion/distributed/test_tensor_parallel.py
+        - >-
+          timeout 20m pytest -s -v
+          tests/diffusion/attention/test_ulysses_uaa.py::test_ulysses_uaa_2d_mask_layout_matches_baseline
+          -m 'core_model and cards_2'
+          --run-level "core_model"
 
     - label: "Diffusion Model Test"
       agent_pool: mi300_1

.buildkite/amd/test-amd-ready.yml

@@ -41,7 +41,9 @@ steps:
         - export VLLM_ROCM_USE_AITER=0
         # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
         # that is not supported on CPU
-        - "pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
+        # Keep multi-GPU tests out of this one-GPU pool even when their module
+        # also carries the cpu marker for unit tests.
+        - "pytest -sv tests/diffusion -m 'core_model and cpu and not (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)' --ignore=tests/diffusion/cache/test_teacache_extractors.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
 
     - label: "Simple · Engine&Entrypoints Test"
       agent_pool: mi300_1
@@ -141,6 +143,11 @@ steps:
         - export VLLM_LOGGING_LEVEL=DEBUG
         - export VLLM_WORKER_MULTIPROC_METHOD=spawn
         - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
+        - >-
+          timeout 20m pytest -s -v
+          tests/diffusion/attention/test_ulysses_uaa.py::test_ulysses_uaa_2d_mask_layout_matches_baseline
+          -m 'core_model and cards_2'
+          --run-level "core_model"
         - >-
           timeout 20m pytest -s -v
           tests/diffusion/models/ltx2/test_ltx2_transformer_ulysses.py

tests/buildkite/test_amd_pipeline.py

@@ -14,6 +14,11 @@
 AMD_READY_PIPELINE = Path(".buildkite/amd/test-amd-ready.yml")
 AMD_TEMPLATE = Path(".buildkite/amd/test-template-amd-omni.j2")
 
+MULTI_GPU_MARKER_EXCLUSION = (
+    "core_model and cpu and not (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)"
+)
+ULYSSES_UAA_2D_NODE = "tests/diffusion/attention/test_ulysses_uaa.py::test_ulysses_uaa_2d_mask_layout_matches_baseline"
+
 
 def _find_step(label: str, pipeline_path: Path = AMD_MERGE_PIPELINE) -> dict:
     pipeline = yaml.safe_load(pipeline_path.read_text(encoding="utf-8"))
@@ -74,6 +79,37 @@ def test_z_image_merge_timeout_covers_cold_aiter_compile() -> None:
     assert split(pytest_command)[:2] == ["timeout", "55m"]
 
 
+@pytest.mark.parametrize("pipeline_path", [AMD_READY_PIPELINE, AMD_MERGE_PIPELINE])
+def test_diffusion_cpu_suite_excludes_multi_gpu_tests(pipeline_path: Path) -> None:
+    step = _find_step("Simple · Diffusion Test · Shard %N/%t", pipeline_path)
+    pytest_command = next(command for command in step["commands"] if "pytest" in command)
+    argv = split(pytest_command)
+
+    marker_index = argv.index("-m")
+    assert argv[marker_index + 1] == MULTI_GPU_MARKER_EXCLUSION
+
+
+@pytest.mark.parametrize(
+    ("pipeline_path", "step_label"),
+    [
+        (AMD_READY_PIPELINE, "Diffusion Sequence Parallelism Test"),
+        (AMD_MERGE_PIPELINE, "Diffusion Tensor Parallelism Test"),
+    ],
+)
+def test_ulysses_uaa_2d_mask_runs_on_two_gpu_lane(
+    pipeline_path: Path,
+    step_label: str,
+) -> None:
+    step = _find_step(step_label, pipeline_path)
+    pytest_command = next(command for command in step["commands"] if ULYSSES_UAA_2D_NODE in command)
+    argv = split(pytest_command)
+
+    assert step["agent_pool"] == "mi300_2"
+    assert ULYSSES_UAA_2D_NODE in argv
+    marker_index = argv.index("-m")
+    assert argv[marker_index + 1] == "core_model and cards_2"
+
+
 def test_cosyvoice_ready_smoke_uses_sdpa() -> None:
     step = _find_step("CosyVoice3-TTS E2E Smoke (SDPA)", AMD_READY_PIPELINE)