From b9cb57b6310de8bbc85a278e165ddf0690a5667c Mon Sep 17 00:00:00 2001
From: TJian <tunjian.tan@embeddedllm.com>
Date: Wed, 13 May 2026 20:50:57 +0800
Subject: [PATCH] [ROCm] Bugfix wan22 (#3463)

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
---
 .buildkite/bootstrap-amd-omni.sh   | 20 ++++++++++++++++--
 .buildkite/test-amd-merge.yml      | 12 +++++++++++
 .buildkite/test-amd-ready.yaml     | 21 +++++++++++++++++++
 vllm_omni/diffusion/layers/rope.py | 33 +++++++++++++++++++++++++++++-
 4 files changed, 83 insertions(+), 3 deletions(-)

diff --git a/.buildkite/bootstrap-amd-omni.sh b/.buildkite/bootstrap-amd-omni.sh
index 284fd371a..16ed02b26 100755
--- a/.buildkite/bootstrap-amd-omni.sh
+++ b/.buildkite/bootstrap-amd-omni.sh
@@ -173,8 +173,24 @@ upload_pipeline() {
 
     cd .buildkite
 
-    # Select test definition file: merge suite for main, ready suite for PRs
-    if [[ $BUILDKITE_BRANCH == "main" ]]; then
+    # Select test definition file: merge suite for main, ready suite for PRs.
+    # For debugging, DEBUG_TEST_YAML can override the selection — accepts
+    # "merge" or "ready" (case-insensitive).
+    if [[ -n "${DEBUG_TEST_YAML:-}" ]]; then
+        case "${DEBUG_TEST_YAML,,}" in
+            merge)
+                TEST_YAML="test-amd-merge.yml"
+                ;;
+            ready)
+                TEST_YAML="test-amd-ready.yaml"
+                ;;
+            *)
+                echo "ERROR: DEBUG_TEST_YAML must be 'merge' or 'ready', got '$DEBUG_TEST_YAML'" >&2
+                exit 1
+                ;;
+        esac
+        echo "DEBUG_TEST_YAML override: using $TEST_YAML"
+    elif [[ $BUILDKITE_BRANCH == "main" ]]; then
         TEST_YAML="test-amd-merge.yml"
     else
         TEST_YAML="test-amd-ready.yaml"
diff --git a/.buildkite/test-amd-merge.yml b/.buildkite/test-amd-merge.yml
index 5c76ed742..2332e9f01 100644
--- a/.buildkite/test-amd-merge.yml
+++ b/.buildkite/test-amd-merge.yml
@@ -193,6 +193,18 @@ steps:
 #     - export VLLM_ROCM_USE_AITER_RMSNORM=0
 #     - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
 
+- label: "Diffusion · Wan22 Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - |
+      timeout 40m bash -c "
+        export VLLM_IMAGE_FETCH_TIMEOUT=60
+        pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
+      "
+
 - label: "Voxtral-TTS E2E Test"
   agent_pool: mi325_1
   depends_on: amd-build
diff --git a/.buildkite/test-amd-ready.yaml b/.buildkite/test-amd-ready.yaml
index c608b3837..c43d8a90f 100644
--- a/.buildkite/test-amd-ready.yaml
+++ b/.buildkite/test-amd-ready.yaml
@@ -94,6 +94,27 @@ steps:
   commands:
     - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
 
+- label: "Diffusion · Model Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model"
+
+- label: "Diffusion · Wan22 Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - |
+      timeout 40m bash -c "
+        export VLLM_TEST_CLEAN_GPU_MEMORY=1
+        export VLLM_IMAGE_FETCH_TIMEOUT=60
+        pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model' --run-level 'core_model'
+      "
+
 - label: "Engine Test"
   agent_pool: mi325_1
   depends_on: amd-build
diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py
index 61b62b2f4..b5c4f9df9 100644
--- a/vllm_omni/diffusion/layers/rope.py
+++ b/vllm_omni/diffusion/layers/rope.py
@@ -204,7 +204,38 @@ class RotaryEmbeddingWan(RotaryEmbedding):
         cos: torch.Tensor,
         sin: torch.Tensor,
     ) -> torch.Tensor:
-        return self.forward_native(x, cos, sin)
+        from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
+
+        if cos.dim() > 2:
+            cos = cos.reshape(-1, cos.shape[-1])
+            sin = sin.reshape(-1, sin.shape[-1])
+
+        return apply_rotary_emb(
+            x,
+            cos,
+            sin,
+            interleaved=self.interleaved,
+        )
+
+    def forward_hip(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        if self.apply_rotary_emb_flash_attn is None:
+            return self.forward_native(x, cos, sin)
+
+        if cos.dim() > 2:
+            cos = cos.reshape(-1, cos.shape[-1])
+            sin = sin.reshape(-1, sin.shape[-1])
+
+        return self.apply_rotary_emb_flash_attn(
+            x,
+            cos,
+            sin,
+            interleaved=self.interleaved,
+        )
 
     def forward_npu(
         self,