From 2aa0db7d9cfadc5f0a596328a641a6f09d5887e3 Mon Sep 17 00:00:00 2001
From: Makcum888e <79456407+Makcum888e@users.noreply.github.com>
Date: Tue, 17 Feb 2026 16:54:19 +0300
Subject: [PATCH] [Diffusion] [NPU] Fix CI run (#18921)

---
 .github/workflows/pr-test-npu.yml                           | 6 +++---
 python/sglang/multimodal_gen/runtime/managers/gpu_worker.py | 4 ++--
 .../test/server/ascend/perf_baselines_npu.json              | 6 +++---
 .../test/server/ascend/testcase_configs_npu.py              | 2 +-
 4 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/pr-test-npu.yml b/.github/workflows/pr-test-npu.yml
index e0bcf3731..7169f4718 100644
--- a/.github/workflows/pr-test-npu.yml
+++ b/.github/workflows/pr-test-npu.yml
@@ -28,9 +28,9 @@ jobs:
   check-changes:
     runs-on: ubuntu-latest
     outputs:
-      changes_exist: ${{ steps.filter.outputs.main_package || steps.filter.outputs.multimodal_gen || steps.run-mode.outputs.run_all_tests}}
-      main_package: ${{ steps.filter.outputs.main_package || steps.run-mode.outputs.run_all_tests }}
-      multimodal_gen: ${{ steps.filter.outputs.multimodal_gen || steps.run-mode.outputs.run_all_tests }}
+      changes_exist: ${{ steps.filter.outputs.main_package == 'true' || steps.filter.outputs.multimodal_gen == 'true' || steps.run-mode.outputs.run_all_tests == 'true'}}
+      main_package: ${{ steps.filter.outputs.main_package == 'true' || steps.run-mode.outputs.run_all_tests == 'true' }}
+      multimodal_gen: ${{ steps.filter.outputs.multimodal_gen == 'true' || steps.run-mode.outputs.run_all_tests == 'true' }}
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
index bd92efbd2..ef04a63c2 100644
--- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
+++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
@@ -154,8 +154,8 @@ class GPUWorker:
             output_batch.timings.record_memory_snapshot("mem_analysis", final_snapshot)
 
         # for details on max_memory_reserved: https://docs.pytorch.org/docs/stable/generated/torch.cuda.memory.max_memory_reserved.html
-        peak_reserved_bytes = torch.cuda.max_memory_reserved()
-        peak_allocated_bytes = torch.cuda.max_memory_allocated()
+        peak_reserved_bytes = torch.get_device_module().max_memory_reserved()
+        peak_allocated_bytes = torch.get_device_module().max_memory_allocated()
 
         output_batch.peak_memory_mb = peak_reserved_bytes / (1024**2)
         peak_reserved_gb = peak_reserved_bytes / (1024**3)
diff --git a/python/sglang/multimodal_gen/test/server/ascend/perf_baselines_npu.json b/python/sglang/multimodal_gen/test/server/ascend/perf_baselines_npu.json
index 12d05de79..81a6f32b0 100644
--- a/python/sglang/multimodal_gen/test/server/ascend/perf_baselines_npu.json
+++ b/python/sglang/multimodal_gen/test/server/ascend/perf_baselines_npu.json
@@ -17,9 +17,9 @@
                 "per_frame_generation": null
             },
             "denoise_step_ms": {
-                "0": 195.27,
-                "1": 329.05,
-                "2": 545.43,
+                "0": 364.97,
+                "1": 542.66,
+                "2": 542.1,
                 "3": 541.3,
                 "4": 537.07,
                 "5": 537.21,
diff --git a/python/sglang/multimodal_gen/test/server/ascend/testcase_configs_npu.py b/python/sglang/multimodal_gen/test/server/ascend/testcase_configs_npu.py
index 96211ecb9..5a3710733 100644
--- a/python/sglang/multimodal_gen/test/server/ascend/testcase_configs_npu.py
+++ b/python/sglang/multimodal_gen/test/server/ascend/testcase_configs_npu.py
@@ -12,7 +12,7 @@ ONE_NPU_CASES: list[DiffusionTestCase] = [
         DiffusionServerArgs(
             model_path="/root/.cache/modelscope/hub/models/Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
             modality="video",
-            warmup=0,
+            warmup=True,
             custom_validator="video",
         ),
         DiffusionSamplingParams(