[NPU] Enable consistency checking for diffusion tests (#27283)

Co-authored-by: Elizaveta Martirosian <elizaveta.martirosian@gmail.com>
Co-authored-by: Elizaveta Martirosian <you@example.com>

python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py

@@ -128,6 +128,15 @@ def apply_flashinfer_rope_qk_inplace(
             cos = cos_sin_cache[positions, :half_size].to(q.dtype)
             sin = cos_sin_cache[positions, half_size:].to(q.dtype)
 
+        if current_platform.is_npu():
+            q_flat = q.reshape(bsz * seqlen, q_heads, d)
+            k_flat = k.reshape(bsz * seqlen, k_heads, d)
+            q_rot = apply_rotary_embedding(q_flat, cos, sin, interleaved=not is_neox)
+            k_rot = apply_rotary_embedding(k_flat, cos, sin, interleaved=not is_neox)
+            return q_rot.view(bsz, seqlen, q_heads, d), k_rot.view(
+                bsz, seqlen, k_heads, d
+            )
+
         def apply_rope_prefix(x: torch.Tensor, num_heads: int) -> torch.Tensor:
             x_flat = x.reshape(bsz * seqlen, num_heads, d)
             x_rot = x_flat[..., :rope_dim]

python/sglang/multimodal_gen/test/server/ascend/testcase_configs_npu.py

@@ -35,7 +35,6 @@ ONE_NPU_CASES: list[DiffusionTestCase] = [
             extras=EXTRAS_DISABLE_WARMUP,
         ),
         T2I_sampling_params,
-        run_consistency_check=False,
     ),
     # === Text to Video (T2V) ===
     DiffusionTestCase(
@@ -47,7 +46,6 @@ ONE_NPU_CASES: list[DiffusionTestCase] = [
         DiffusionSamplingParams(
             prompt=T2V_PROMPT,
         ),
-        run_consistency_check=False,
     ),
 ]
 
@@ -62,7 +60,6 @@ TWO_NPU_CASES: list[DiffusionTestCase] = [
             extras=EXTRAS_DISABLE_WARMUP,
         ),
         T2I_sampling_params,
-        run_consistency_check=False,
     ),
     DiffusionTestCase(
         "qwen_image_t2i_2npu",
@@ -75,7 +72,6 @@ TWO_NPU_CASES: list[DiffusionTestCase] = [
             extras=EXTRAS_DISABLE_WARMUP,
         ),
         T2I_sampling_params,
-        run_consistency_check=False,
     ),
     # === Text to Video (T2V) ===
     DiffusionTestCase(
@@ -90,7 +86,6 @@ TWO_NPU_CASES: list[DiffusionTestCase] = [
         DiffusionSamplingParams(
             prompt=T2V_PROMPT,
         ),
-        run_consistency_check=False,
     ),
 ]
 

python/sglang/multimodal_gen/test/test_utils.py

@@ -21,6 +21,7 @@ import numpy as np
 import requests
 from PIL import Image, ImageDraw, ImageFont
 
+from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.runtime.utils.perf_logger import (
@@ -34,6 +35,10 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 SGL_TEST_FILES_CI_DATA_REVISION = "caa56302ccf2d289e4488ed06d952edf5d2314cf"
+
+if current_platform.is_npu():
+    SGL_TEST_FILES_CI_DATA_REVISION = "670d66a8a290b62c0c3c077b3e9b0f4a4d9a44e7"
+
 SGL_TEST_FILES_CONSISTENCY_GT_ROOT = (
     "https://raw.githubusercontent.com/"
     f"sgl-project/ci-data/{SGL_TEST_FILES_CI_DATA_REVISION}/"
@@ -51,7 +56,14 @@ SGL_TEST_FILES_OFFICIAL_CONSISTENCY_GT_BASE_ASCEND = (
 SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE_ASCEND = (
     f"{SGL_TEST_FILES_CONSISTENCY_GT_ROOT}/sglang_generated/ascend"
 )
+
 SGL_TEST_FILES_CONSISTENCY_GT_BASE = SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE
+
+if current_platform.is_npu():
+    SGL_TEST_FILES_CONSISTENCY_GT_BASE = (
+        SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE_ASCEND
+    )
+
 SGL_TEST_FILES_CONSISTENCY_GT_BASES = (
     SGL_TEST_FILES_OFFICIAL_CONSISTENCY_GT_BASE,
     SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE,
@@ -1018,7 +1030,7 @@ def _find_remote_consistency_gt_files(
         bases = SGL_TEST_FILES_CONSISTENCY_GT_BASES
     else:
         # Avoid accidentally comparing non-comparable CI cases against official GT.
-        bases = (SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE,)
+        bases = (SGL_TEST_FILES_CONSISTENCY_GT_BASE,)
     for base_url in bases:
         candidates = _remote_consistency_gt_candidates(
             base_url, case_id, num_gpus, is_video, output_format