[CI][Test] Add NPU nightly tests (#3480)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>

.buildkite/pipeline-npu-a3.yaml

@@ -0,0 +1,110 @@
+# steps:
+#   - label: ":buildkit: Build and Push NPU Test Image - A3"
+#     key: image-build-a3
+#     agents:
+#       queue: "ascend-a3"
+#       resource_class: "npu-2"
+#     plugins:
+#       - kubernetes:
+#           metadata:
+#             annotations:
+#               vault.hashicorp.com/agent-init-first: "true"
+#               vault.hashicorp.com/agent-inject: "true"
+#               vault.hashicorp.com/agent-inject-perms-ca.pem: "0400"
+#               vault.hashicorp.com/agent-inject-perms-cert.pem: "0400"
+#               vault.hashicorp.com/agent-inject-perms-config.json: "0400"
+#               vault.hashicorp.com/agent-inject-perms-key.pem: "0400"
+#               vault.hashicorp.com/agent-inject-secret-ca.pem: internal/data/ascend/buildkitd
+#               vault.hashicorp.com/agent-inject-secret-cert.pem: internal/data/ascend/buildkitd
+#               vault.hashicorp.com/agent-inject-secret-config.json: internal/data/ascend/buildkitd
+#               vault.hashicorp.com/agent-inject-secret-key.pem: internal/data/ascend/buildkitd
+#               vault.hashicorp.com/agent-inject-template-ca.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.RootCA }}\n{{- end }}"
+#               vault.hashicorp.com/agent-inject-template-cert.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaCert }}\n{{- end }}"
+#               vault.hashicorp.com/agent-inject-template-config.json: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.dockerConfig }}\n{{- end }}"
+#               vault.hashicorp.com/agent-inject-template-key.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaKey }}\n{{- end }}"
+#               vault.hashicorp.com/agent-pre-populate-only: "true"
+#               vault.hashicorp.com/agent-run-as-group: "1000"
+#               vault.hashicorp.com/agent-run-as-user: "1000"
+#               vault.hashicorp.com/agent-service-account-token-volume-name: token-vol
+#               vault.hashicorp.com/role: ascend-gha-runners
+#               vault.hashicorp.com/secret-volume-path: /home/user/.docker/
+#               vault.hashicorp.com/tls-skip-verify: "true"
+#           podSpecPatch:
+#             volumes:
+#             - name: token-vol
+#               projected:
+#                 defaultMode: 420
+#                 sources:
+#                 - serviceAccountToken:
+#                     audience: api
+#                     expirationSeconds: 600
+#                     path: token
+#     env:
+#       VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
+#       IMAGE_NAME: "vllm-omni-ci-npu"
+#       IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry"
+#       BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234"
+#     command: |
+#       set -ex
+
+#       echo "--- Building and pushing NPU Test Image"
+#       echo "Image: $${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
+#       echo "buildkitd address: $${BUILDKITD_ADDR}"
+
+#       git config --global url."https://gh-proxy.test.osinfra.cn/https://github.com/".insteadOf "https://github.com/"
+
+#       if ! command -v buildctl &> /dev/null; then
+#         echo "Installing buildctl..."
+#         mkdir -p /tmp/buildkit
+#         BUILDKIT_VERSION="v0.29.0"
+#         wget -q "https://gh-proxy.test.osinfra.cn/https://github.com/moby/buildkit/releases/download/$${BUILDKIT_VERSION}/buildkit-$${BUILDKIT_VERSION}.linux-arm64.tar.gz" -O /tmp/buildkit.tar.gz
+#         tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit
+#         cp /tmp/buildkit/bin/buildctl /usr/local/bin/
+#       fi
+
+#       export DOCKER_CONFIG=/home/user/.docker
+#       buildctl \
+#         --addr="$${BUILDKITD_ADDR}" \
+#         --tlscacert=/home/user/.docker/ca.pem \
+#         --tlscert=/home/user/.docker/cert.pem \
+#         --tlskey=/home/user/.docker/key.pem \
+#         build \
+#         --frontend dockerfile.v0 \
+#         --local context=. \
+#         --local dockerfile=./docker \
+#         --opt filename=Dockerfile.npu.ci.a3 \
+#         --secret id=dockerconfig,src=/home/user/.docker/config.json \
+#         --output type=image,name=$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG},push=true \
+#         --progress=plain
+
+#       echo "--- Image pushed successfully"
+#       echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
+
+#   # - label: "Upload Ready Pipeline - A3"
+#   #   depends_on: image-build-a3
+#   #   key: upload-ready-pipeline-a3
+#   #   if: build.branch != "main" && build.pull_request.labels includes "npu-test"
+#   #   agents:
+#   #     queue: "ascend-a3"
+#   #     resource_class: "npu-2"
+#   #   command: |
+#   #     buildkite-agent pipeline upload .buildkite/test-npu-ready.yaml
+
+  # # L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild)
+  # - label: "Upload Nightly Pipeline"
+  #   depends_on: image-build
+  #   key: upload-nightly-pipeline
+  #   if: >-
+  #     (build.branch == "main" && build.env("NIGHTLY") == "1") ||
+  #     (build.branch != "main" && (
+  #       build.pull_request.labels includes "nightly-test" ||
+  #       build.pull_request.labels includes "omni-test" ||
+  #       build.pull_request.labels includes "tts-test" ||
+  #       build.pull_request.labels includes "diffusion-x2iat-test" ||
+  #       build.pull_request.labels includes "diffusion-x2v-test"
+  #     ))
+  #   commands:
+  #     - buildkite-agent pipeline upload .buildkite/test-npu-nightly.yaml
+  #   agents:
+  #     queue: "ascend-a3"
+  #     resource_class: "npu-2"

.buildkite/pipeline-npu.yaml

@@ -0,0 +1,108 @@
+steps:
+  - label: ":buildkit: Build and Push NPU Test Image - A2"
+    key: image-build
+    agents:
+      queue: "ascend-a2b3"
+      resource_class: "npu-1"
+    plugins:
+      - kubernetes:
+          metadata:
+            annotations:
+              vault.hashicorp.com/agent-init-first: "true"
+              vault.hashicorp.com/agent-inject: "true"
+              vault.hashicorp.com/agent-inject-perms-ca.pem: "0400"
+              vault.hashicorp.com/agent-inject-perms-cert.pem: "0400"
+              vault.hashicorp.com/agent-inject-perms-config.json: "0400"
+              vault.hashicorp.com/agent-inject-perms-key.pem: "0400"
+              vault.hashicorp.com/agent-inject-secret-ca.pem: internal/data/ascend/buildkitd
+              vault.hashicorp.com/agent-inject-secret-cert.pem: internal/data/ascend/buildkitd
+              vault.hashicorp.com/agent-inject-secret-config.json: internal/data/ascend/buildkitd
+              vault.hashicorp.com/agent-inject-secret-key.pem: internal/data/ascend/buildkitd
+              vault.hashicorp.com/agent-inject-template-ca.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.RootCA }}\n{{- end }}"
+              vault.hashicorp.com/agent-inject-template-cert.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaCert }}\n{{- end }}"
+              vault.hashicorp.com/agent-inject-template-config.json: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.dockerConfig }}\n{{- end }}"
+              vault.hashicorp.com/agent-inject-template-key.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaKey }}\n{{- end }}"
+              vault.hashicorp.com/agent-pre-populate-only: "true"
+              vault.hashicorp.com/agent-run-as-group: "1000"
+              vault.hashicorp.com/agent-run-as-user: "1000"
+              vault.hashicorp.com/agent-service-account-token-volume-name: token-vol
+              vault.hashicorp.com/role: ascend-gha-runners
+              vault.hashicorp.com/secret-volume-path: /home/user/.docker/
+              vault.hashicorp.com/tls-skip-verify: "true"
+          podSpecPatch:
+            volumes:
+            - name: token-vol
+              projected:
+                defaultMode: 420
+                sources:
+                - serviceAccountToken:
+                    audience: api
+                    expirationSeconds: 600
+                    path: token
+    env:
+      VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
+      IMAGE_NAME: "vllm-omni-ci-npu"
+      IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry"
+      BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234"
+    command: |
+      set -ex
+
+      echo "--- Building and pushing NPU Test Image"
+      echo "Image: $${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
+      echo "buildkitd address: $${BUILDKITD_ADDR}"
+
+      if ! command -v buildctl &> /dev/null; then
+        echo "Installing buildctl..."
+        mkdir -p /tmp/buildkit
+        BUILDKIT_VERSION="v0.29.0"
+        wget -q "https://github.com/moby/buildkit/releases/download/$${BUILDKIT_VERSION}/buildkit-$${BUILDKIT_VERSION}.linux-arm64.tar.gz" -O /tmp/buildkit.tar.gz
+        tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit
+        cp /tmp/buildkit/bin/buildctl /usr/local/bin/
+      fi
+
+      export DOCKER_CONFIG=/home/user/.docker
+      buildctl \
+        --addr="$${BUILDKITD_ADDR}" \
+        --tlscacert=/home/user/.docker/ca.pem \
+        --tlscert=/home/user/.docker/cert.pem \
+        --tlskey=/home/user/.docker/key.pem \
+        build \
+        --frontend dockerfile.v0 \
+        --local context=. \
+        --local dockerfile=./docker \
+        --opt filename=Dockerfile.npu.ci \
+        --secret id=dockerconfig,src=/home/user/.docker/config.json \
+        --output type=image,name=$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG},push=true \
+        --progress=plain
+
+      echo "--- Image pushed successfully"
+      echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
+
+  # - label: "Upload Ready Pipeline"
+  #   depends_on: image-build
+  #   key: upload-ready-pipeline
+  #   if: build.branch != "main" && build.pull_request.labels includes "npu-test"
+  #   agents:
+  #     queue: "ascend-a2b3"
+  #     resource_class: "npu-1"
+  #   command: |
+  #     buildkite-agent pipeline upload .buildkite/test-npu-ready.yaml
+
+  # L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild)
+  - label: "Upload Nightly Pipeline"
+    depends_on: image-build
+    key: upload-nightly-pipeline
+    if: >-
+      (build.branch == "main" && build.env("NIGHTLY") == "1") ||
+      (build.branch != "main" && (
+        build.pull_request.labels includes "nightly-test" ||
+        build.pull_request.labels includes "omni-test" ||
+        build.pull_request.labels includes "tts-test" ||
+        build.pull_request.labels includes "diffusion-x2iat-test" ||
+        build.pull_request.labels includes "diffusion-x2v-test"
+      ))
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-npu-nightly.yaml
+    agents:
+      queue: "ascend-a2b3"
+      resource_class: "npu-1"

.buildkite/scripts/upload_pipeline_with_skip_ci.sh

@@ -1,10 +1,18 @@
 #!/usr/bin/env bash
-# Evaluate docs-only skip-ci and upload continuation steps from the same `.buildkite/pipeline.yml`
+# Evaluate docs-only skip-ci and upload continuation steps from the given pipeline YAML
 # (YAML document after the first `---`). Buildkite `if` is evaluated at upload time.
+#
+# Usage: upload_pipeline_with_skip_ci.sh [pipeline_yaml]
+#   pipeline_yaml: path relative to repo root or absolute (default: .buildkite/pipeline.yml)
 set -euo pipefail
 
 ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
-PIPELINE_YML="${ROOT}/.buildkite/pipeline.yml"
+PIPELINE_ARG="${1:-.buildkite/pipeline.yml}"
+if [[ "${PIPELINE_ARG}" = /* ]]; then
+  PIPELINE_YML="${PIPELINE_ARG}"
+else
+  PIPELINE_YML="${ROOT}/${PIPELINE_ARG}"
+fi
 
 # Prints a single digit to stdout: 1 = skip image CI, 0 = run. Logs go to stderr.
 is_docs_only_change() {
@@ -104,12 +112,13 @@ import pathlib
 path = pathlib.Path(os.environ["PIPELINE_YML"])
 text = path.read_text(encoding="utf-8")
 sep = "\n---\n"
-if sep not in text:
-    raise SystemExit(
-        "upload_pipeline_with_skip_ci: .buildkite/pipeline.yml must contain a '\\n---\\n' separator "
-        "(document 1 = bootstrap, document 2 = uploaded steps)"
-    )
-_, continuation = text.split(sep, 1)
+# Two supported layouts:
+#   - multi-doc: doc 1 = bootstrap (loaded by Buildkite), doc 2 (after `---`) = uploaded steps.
+#   - single-doc: whole file is the uploaded steps; caller (UI init step) is the bootstrap.
+if sep in text:
+    _, continuation = text.split(sep, 1)
+else:
+    continuation = text
 
 skip = os.environ.get("SKIP_CI") == "1"
 # When docs-only skip-ci: skip default CI image, but still build for L4 nightly (PR label nightly-test or

.buildkite/test-npu-nightly.yaml

@@ -0,0 +1,24 @@
+steps:
+  - group: ":card_index_dividers: Diffusion X2V Model Test"
+    key: nightly-diffusion-x2v-group
+    depends_on: upload-nightly-pipeline
+    if: >-
+      build.env("NIGHTLY") == "1" ||
+      build.pull_request.labels includes "nightly-test" ||
+      build.pull_request.labels includes "diffusion-x2v-test"
+    steps:
+      - label: ":full_moon: Diffusion X2V · Function Test"
+        agents:
+          queue: "ascend-a2b3"
+          resource_class: "npu-4"
+        image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}"
+        plugins:
+          - kubernetes:
+              podSpecPatch:
+                imagePullSecrets:
+                  - name: swr-secret
+        env:
+          VLLM_WORKER_MULTIPROC_METHOD: spawn
+          HF_TOKEN: "${HF_TOKEN}"
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k npu_i2v --run-level "full_model"

docker/Dockerfile.npu.ci

@@ -0,0 +1,42 @@
+ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
+ARG VLLM_ASCEND_TAG=nightly-main
+FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
+
+ARG APP_DIR=/vllm-workspace/vllm-omni
+WORKDIR ${APP_DIR}
+COPY . .
+
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+        ffmpeg \
+        libgl1 \
+        libglib2.0-0 \
+        libgomp1 \
+        libsm6 \
+        libxext6 \
+        libxrender1 \
+        libxcb1 && \
+    rm -rf /var/cache/apt/* && \
+    rm -rf /var/lib/apt/lists/*
+
+RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    python3 -m pip install -U pip && \
+    python3 -m pip install -v \
+      "pytest>=7.0.0" \
+      "pytest-asyncio>=0.21.0" \
+      "pytest-cov>=4.0.0" \
+      "pytest-mock>=3.10.0" \
+      "openpyxl>=3.0.0" \
+      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
+      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
+      --default-timeout=60 && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation \
+      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
+      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
+      --default-timeout=60
+
+ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+
+ENTRYPOINT []

docker/Dockerfile.npu.ci.a3

@@ -0,0 +1,42 @@
+ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
+ARG VLLM_ASCEND_TAG=nightly-main-a3
+FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
+
+ARG APP_DIR=/vllm-workspace/vllm-omni
+WORKDIR ${APP_DIR}
+COPY . .
+
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+        ffmpeg \
+        libgl1 \
+        libglib2.0-0 \
+        libgomp1 \
+        libsm6 \
+        libxext6 \
+        libxrender1 \
+        libxcb1 && \
+    rm -rf /var/cache/apt/* && \
+    rm -rf /var/lib/apt/lists/*
+
+RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    python3 -m pip install -U pip && \
+    python3 -m pip install -v \
+      "pytest>=7.0.0" \
+      "pytest-asyncio>=0.21.0" \
+      "pytest-cov>=4.0.0" \
+      "pytest-mock>=3.10.0" \
+      "openpyxl>=3.0.0" \
+      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
+      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
+      --default-timeout=60 && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation \
+      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
+      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
+      --default-timeout=60
+
+ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+
+ENTRYPOINT []

tests/e2e/online_serving/test_wan22_expansion.py

@@ -1,18 +1,16 @@
 """
-Comprehensive tests of diffusion features that are available in online serving mode
-and are supported by the following models:
+Comprehensive tests of diffusion features that are available in online serving mode.
+
+CUDA coverage (3 models × 6 features):
 - Wan-AI/Wan2.2-T2V-A14B-Diffusers
 - Wan-AI/Wan2.2-I2V-A14B-Diffusers
 - Wan-AI/Wan2.2-TI2V-5B-Diffusers
+Features: Cache-DiT, CFG-Parallel, Ulysses-SP, Tensor-Parallel + VAE-Patch-Parallel,
+HSDP, Ring-Attn.
 
-Coverage:
-- Cache-DiT
-- CFG-Parallel
-- Ulysses-SP
-- Tensor-Parallel
-- VAE-Patch-Parallel
-- HSDP
-- Ring-Attn
+NPU coverage (Wan-AI/Wan2.2-I2V-A14B-Diffusers only): 2 cases.
+- 4-card combined: cfg=2 + usp=2 + vae-patch=2 + hsdp.
+- 2-card tp_layerwise: tp=2 + enable-layerwise-offload.
 
 assert_diffusion_response validates successful generation
 """
@@ -27,14 +25,24 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
 PROMPT = "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
 NEGATIVE_PROMPT = "low quality, blurry, distorted face, extra limbs, bad anatomy, watermark, logo, text, ugly, deformed, mutated, jpeg artifacts"
-SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
-PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+
+# CUDA marks (original matrix, unchanged)
+CUDA_SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
+CUDA_PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+
+# NPU marks
+NPU_TWO_CARD_MARKS = hardware_marks(res={"npu": "A2"}, num_cards=2)
+NPU_FOUR_CARD_MARKS = hardware_marks(res={"npu": "A2"}, num_cards=4)
 
 WAN22_MODELS = [
     ("Wan-AI/Wan2.2-T2V-A14B-Diffusers", "t2v"),
     ("Wan-AI/Wan2.2-I2V-A14B-Diffusers", "i2v"),
     ("Wan-AI/Wan2.2-TI2V-5B-Diffusers", "ti2v"),
 ]
+NPU_MODELS = [("Wan-AI/Wan2.2-I2V-A14B-Diffusers", "i2v")]
+
+CACHE_DIT_ARGS = ["--cache-backend", "cache_dit", "--enable-layerwise-offload"]
+HSDP_ARGS = ["--use-hsdp", "--hsdp-shard-size", "2"]
 
 PARALLEL_CONFIGS = [
     ("cfg_parallel", ["--cfg-parallel-size", "2"]),
@@ -44,37 +52,66 @@ PARALLEL_CONFIGS = [
     ("ring_atten", ["--ring", "2"]),
 ]
 
+# NPU: 2 cases only.
+NPU_PARALLEL_CONFIGS = [
+    (
+        "combined",
+        [
+            "--cfg-parallel-size",
+            "2",
+            "--usp",
+            "2",
+            "--vae-patch-parallel-size",
+            "4",
+            "--use-hsdp",
+            "--hsdp-shard-size",
+            "4",
+        ],
+        NPU_FOUR_CARD_MARKS,
+    ),
+    (
+        "tp_layerwise_offload",
+        ["--tensor-parallel-size", "2", "--enable-layerwise-offload"],
+        NPU_TWO_CARD_MARKS,
+    ),
+]
+
 
 def _get_wan22_feature_cases():
     """
-    Generate parameterized test cases covering:
-    - All 3 Wan2.2 model variants with architecture awareness
-    - 1 single-card feature (Cache-DiT)
-    - 6 multi-card parallelism features with CORRECT PARAMETER NAMES per spec
+    Generate parameterized test cases:
+    - CUDA: 3 models × (Cache-DiT + 5 parallel features), original matrix.
+    - NPU: I2V-A14B only, 2 cases (4-card combined, 2-card tp_layerwise_offload).
     """
     cases = []
 
-    # Single-card: Cache-DiT (applies to all models)
+    # ---- CUDA cases (unchanged) ----
     for model_path, model_key in WAN22_MODELS:
         cases.append(
             pytest.param(
-                OmniServerParams(
-                    model=model_path,
-                    server_args=["--cache-backend", "cache_dit", "--enable-layerwise-offload"],
-                ),
-                id=f"{model_key}_cache_dit",
-                marks=SINGLE_CARD_FEATURE_MARKS,
+                OmniServerParams(model=model_path, server_args=CACHE_DIT_ARGS),
+                id=f"cuda_{model_key}_cache_dit",
+                marks=CUDA_SINGLE_CARD_MARKS,
             )
         )
-
-    # Multi-card features
     for model_path, model_key in WAN22_MODELS:
         for feat_id, server_args in PARALLEL_CONFIGS:
             cases.append(
                 pytest.param(
                     OmniServerParams(model=model_path, server_args=server_args),
-                    id=f"{model_key}_{feat_id}",
-                    marks=PARALLEL_FEATURE_MARKS,
+                    id=f"cuda_{model_key}_{feat_id}",
+                    marks=CUDA_PARALLEL_MARKS,
+                )
+            )
+
+    # ---- NPU cases (I2V-A14B only) ----
+    for model_path, model_key in NPU_MODELS:
+        for feat_id, server_args, marks in NPU_PARALLEL_CONFIGS:
+            cases.append(
+                pytest.param(
+                    OmniServerParams(model=model_path, server_args=server_args),
+                    id=f"npu_{model_key}_{feat_id}",
+                    marks=marks,
                 )
             )
 

tests/helpers/env.py

@@ -1,7 +1,7 @@
-"""Test environment / lifecycle helpers (GPU cleanup hooks and memory monitoring for tests).
+"""Test environment / lifecycle helpers (device cleanup hooks and memory monitoring for tests).
 
-``vllm.platforms`` / ``vllm_omni.platforms`` are imported only inside functions that need them
-so importing this module at pytest plugin load does not run before session autouse fixtures
+``vllm_omni.platforms`` is imported only inside functions that need it so importing this module
+at pytest plugin load does not run before session autouse fixtures.
 """
 
 from __future__ import annotations
@@ -15,6 +15,8 @@ from contextlib import contextmanager
 
 import torch
 
+from vllm_omni.platforms import current_omni_platform
+
 
 def get_physical_device_indices(devices):
     visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
@@ -32,35 +34,27 @@ def wait_for_gpu_memory_to_clear(
     threshold_ratio: float | None = None,
     timeout_s: float = 120,
 ) -> None:
-    from vllm.platforms import current_platform
-
     assert threshold_bytes is not None or threshold_ratio is not None
     devices = get_physical_device_indices(devices)
     start_time = time.time()
 
     device_list = ", ".join(str(d) for d in devices)
     if threshold_bytes is not None:
-        threshold_str = f"{threshold_bytes / 2**30:.2f} GiB"
-        condition_str = f"Memory usage ≤ {threshold_str}"
-    else:
-        threshold_percent = threshold_ratio * 100
-        threshold_str = f"{threshold_percent:.1f}%"
-        condition_str = f"Memory usage ratio ≤ {threshold_str}"
-
-    print(f"[GPU Memory Monitor] Waiting for GPU {device_list} to free memory, Condition: {condition_str}")
-
-    if threshold_bytes is not None:
+        condition_str = f"Memory usage ≤ {threshold_bytes / 2**30:.2f} GiB"
 
         def is_free(used, total):
             return used <= threshold_bytes / 2**30
     else:
+        condition_str = f"Memory usage ratio ≤ {threshold_ratio * 100:.1f}%"
 
         def is_free(used, total):
             return used / total <= threshold_ratio
 
+    print(f"[GPU Memory Monitor] Waiting for GPU {device_list} to free memory, Condition: {condition_str}")
+
     @contextmanager
-    def nvml_scope():
-        if current_platform.is_rocm():
+    def smi_scope():
+        if current_omni_platform.is_rocm():
             from amdsmi import amdsmi_init, amdsmi_shut_down
 
             amdsmi_init()
@@ -68,7 +62,7 @@ def wait_for_gpu_memory_to_clear(
                 yield
             finally:
                 amdsmi_shut_down()
-        elif current_platform.is_cuda():
+        elif current_omni_platform.is_cuda():
             from vllm.third_party.pynvml import nvmlInit, nvmlShutdown
 
             nvmlInit()
@@ -79,147 +73,86 @@ def wait_for_gpu_memory_to_clear(
         else:
             yield
 
-    is_rocm = current_platform.is_rocm()
-
-    with nvml_scope():
-        if is_rocm:
+    def get_mem_gib(device: int) -> tuple[float, float]:
+        if current_omni_platform.is_rocm():
             from amdsmi import amdsmi_get_gpu_vram_usage, amdsmi_get_processor_handles
-        elif current_platform.is_cuda():
-            from vllm.third_party.pynvml import nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo
 
+            info = amdsmi_get_gpu_vram_usage(amdsmi_get_processor_handles()[device])
+            return info["vram_used"] / 2**10, info["vram_total"] / 2**10
+        if current_omni_platform.is_npu():
+            free_bytes, total_bytes = torch.npu.mem_get_info(device)
+            return (total_bytes - free_bytes) / 2**30, total_bytes / 2**30
+        from vllm.third_party.pynvml import nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo
+
+        info = nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(device))
+        return info.used / 2**30, info.total / 2**30
+
+    with smi_scope():
         while True:
-            output: dict[int, str] = {}
-            output_raw: dict[int, tuple[float, float]] = {}
-            for device in devices:
-                if is_rocm:
-                    dev_handle = amdsmi_get_processor_handles()[device]
-                    mem_info = amdsmi_get_gpu_vram_usage(dev_handle)
-                    gb_used = mem_info["vram_used"] / 2**10
-                    gb_total = mem_info["vram_total"] / 2**10
-                else:
-                    dev_handle = nvmlDeviceGetHandleByIndex(device)
-                    mem_info = nvmlDeviceGetMemoryInfo(dev_handle)
-                    gb_used = mem_info.used / 2**30
-                    gb_total = mem_info.total / 2**30
-                output_raw[device] = (gb_used, gb_total)
-                usage_percent = (gb_used / gb_total) * 100 if gb_total > 0 else 0
-                output[device] = f"{gb_used:.1f}GiB/{gb_total:.1f}GiB ({usage_percent:.1f}%)"
+            output_raw = {d: get_mem_gib(d) for d in devices}
+            output = {
+                d: f"{used:.1f}GiB/{total:.1f}GiB ({(used / total) * 100 if total > 0 else 0:.1f}%)"
+                for d, (used, total) in output_raw.items()
+            }
 
             print("[GPU Memory Status] Current usage:")
             for device_id, mem_info in output.items():
                 print(f"  GPU {device_id}: {mem_info}")
 
             dur_s = time.time() - start_time
-            elapsed_minutes = dur_s / 60
             if all(is_free(used, total) for used, total in output_raw.values()):
                 print(f"[GPU Memory Freed] Devices {device_list} meet memory condition")
                 print(f"   Condition: {condition_str}")
-                print(f"   Wait time: {dur_s:.1f} seconds ({elapsed_minutes:.1f} minutes)")
+                print(f"   Wait time: {dur_s:.1f} seconds ({dur_s / 60:.1f} minutes)")
                 break
 
             if dur_s >= timeout_s:
                 raise ValueError(
                     f"[GPU Memory Timeout] Devices {device_list} still don't meet memory condition after {dur_s:.1f} seconds\n"
                     f"Condition: {condition_str}\n"
-                    f"Current status:\n" + "\n".join(f"  GPU {device}: {output[device]}" for device in devices)
+                    f"Current status:\n" + "\n".join(f"  GPU {d}: {output[d]}" for d in devices)
                 )
 
             gc.collect()
-            torch.accelerator.empty_cache()
+            current_omni_platform.empty_cache()
             time.sleep(5)
 
 
+def _run_smi(label: str, cmd: list[str], head_lines: int, timeout: float = 5) -> None:
+    print("\n" + "=" * 80)
+    print(label)
+    print("=" * 80)
+    try:
+        result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
+        if result.returncode == 0 and result.stdout.strip():
+            lines = result.stdout.strip().split("\n")
+            for line in lines[:head_lines]:
+                print(line)
+            if len(lines) > head_lines:
+                print(f"... (showing first {head_lines} of {len(lines)} lines)")
+        else:
+            print(f"{cmd[0]} command failed or produced no output")
+    except (subprocess.TimeoutExpired, FileNotFoundError):
+        print(f"{cmd[0]} not available or timed out")
+    except Exception as e:
+        print(f"Error running {' '.join(cmd)}: {e}")
+
+
 def _print_gpu_processes() -> None:
-    """Print GPU information including nvidia-smi/amd-smi and system processes."""
-    from vllm.platforms import current_platform
-
-    # Check for NVIDIA GPUs
-    if current_platform.is_cuda():
-        print("\n" + "=" * 80)
-        print("NVIDIA GPU Information (nvidia-smi)")
-        print("=" * 80)
-        try:
-            nvidia_result = subprocess.run(
-                ["nvidia-smi"],
-                capture_output=True,
-                text=True,
-                timeout=5,
-            )
-            if nvidia_result.returncode == 0:
-                lines = nvidia_result.stdout.strip().split("\n")
-                for line in lines[:20]:
-                    print(line)
-                if len(lines) > 20:
-                    print(f"... (showing first 20 of {len(lines)} lines)")
-            else:
-                print("nvidia-smi command failed")
-        except (subprocess.TimeoutExpired, FileNotFoundError):
-            print("nvidia-smi not available or timed out")
-        except Exception as e:
-            print(f"Error running nvidia-smi: {e}")
-
-        print("\n" + "=" * 80)
-        print("Detailed GPU Processes (nvidia-smi pmon)")
-        print("=" * 80)
-        try:
-            pmon_result = subprocess.run(
-                ["nvidia-smi", "pmon", "-c", "1"],
-                capture_output=True,
-                text=True,
-                timeout=3,
-            )
-            if pmon_result.returncode == 0 and pmon_result.stdout.strip():
-                print(pmon_result.stdout)
-            else:
-                print("No active GPU processes found via nvidia-smi pmon")
-        except Exception:
-            print("nvidia-smi pmon not available")
-
-    # Check for AMD ROCm GPUs
-    elif current_platform.is_rocm():
-        print("\n" + "=" * 80)
-        print("AMD GPU Information (amd-smi)")
-        print("=" * 80)
-        try:
-            amd_result = subprocess.run(
-                ["amd-smi"],
-                capture_output=True,
-                text=True,
-                timeout=5,
-            )
-            if amd_result.returncode == 0:
-                lines = amd_result.stdout.strip().split("\n")
-                for line in lines[:30]:  # Show more lines for AMD output
-                    print(line)
-                if len(lines) > 30:
-                    print(f"... (showing first 30 of {len(lines)} lines)")
-            else:
-                print("amd-smi command failed")
-        except (subprocess.TimeoutExpired, FileNotFoundError):
-            print("amd-smi not available or timed out")
-        except Exception as e:
-            print(f"Error running amd-smi: {e}")
-
-        print("\n" + "=" * 80)
-        print("Detailed AMD GPU Processes (amd-smi process)")
-        print("=" * 80)
-        try:
-            process_result = subprocess.run(
-                ["amd-smi", "process"],
-                capture_output=True,
-                text=True,
-                timeout=3,
-            )
-            if process_result.returncode == 0 and process_result.stdout.strip():
-                print(process_result.stdout)
-            else:
-                print("No active GPU processes found via amd-smi process")
-        except Exception:
-            print("amd-smi process not available")
+    """Print device information via nvidia-smi/npu-smi/amd-smi."""
+    from vllm_omni.platforms import current_omni_platform
 
+    if current_omni_platform.is_cuda():
+        _run_smi("NVIDIA GPU Information (nvidia-smi)", ["nvidia-smi"], 20)
+        _run_smi("Detailed GPU Processes (nvidia-smi pmon)", ["nvidia-smi", "pmon", "-c", "1"], 100, timeout=3)
+    elif current_omni_platform.is_npu():
+        _run_smi("Ascend NPU Information (npu-smi info)", ["npu-smi", "info"], 40)
+    elif current_omni_platform.is_rocm():
+        _run_smi("AMD GPU Information (amd-smi)", ["amd-smi"], 30)
+        _run_smi("Detailed AMD GPU Processes (amd-smi process)", ["amd-smi", "process"], 100, timeout=3)
     else:
         print("\n" + "=" * 80)
-        print("WARNING: No supported GPU platform detected (neither CUDA nor ROCm)")
+        print("WARNING: No supported device platform detected")
         print("=" * 80)
 
     print("\n" + "=" * 80)
@@ -230,7 +163,7 @@ def _print_gpu_processes() -> None:
 def run_pre_test_cleanup() -> None:
     print("Pre-test GPU status:")
 
-    num_gpus = torch.accelerator.device_count()
+    num_gpus = current_omni_platform.device_count()
     if num_gpus > 0:
         try:
             wait_for_gpu_memory_to_clear(
@@ -242,9 +175,9 @@ def run_pre_test_cleanup() -> None:
 
 
 def run_post_test_cleanup() -> None:
-    if torch.cuda.is_available():
+    if current_omni_platform.is_available():
         gc.collect()
-        torch.accelerator.empty_cache()
+        current_omni_platform.empty_cache()
 
         print("Post-test GPU status:")
         _print_gpu_processes()

tests/helpers/runtime.py

@@ -4,6 +4,7 @@ import base64
 import concurrent.futures
 import copy
 import errno
+import gc
 import io
 import json
 import os
@@ -26,8 +27,11 @@ import torch
 import yaml
 from openai import APIError, OpenAI, omit
 from PIL import Image
-from vllm import TextPrompt
-from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
+from vllm import TextPrompt, envs
+from vllm.distributed.parallel_state import (
+    destroy_distributed_environment,
+    destroy_model_parallel,
+)
 from vllm.logger import init_logger
 
 from tests.helpers.assertions import (
@@ -49,6 +53,39 @@ from vllm_omni.platforms import current_omni_platform
 logger = init_logger(__name__)
 
 
+def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
+    # Reset environment variable cache
+    envs.disable_envs_cache()
+
+    # Reset rocm_aiter_ops class variables to match current os.environ.
+    # These are class-level attributes that persist across tests and are
+    # NOT restored by monkeypatch (which only restores os.environ).
+    from vllm_omni.platforms import current_omni_platform
+
+    if current_omni_platform.is_rocm():
+        from vllm._aiter_ops import rocm_aiter_ops
+
+        rocm_aiter_ops.refresh_env_variables()
+
+    # Ensure all objects are not frozen before cleanup
+    gc.unfreeze()
+
+    destroy_model_parallel()
+    destroy_distributed_environment()
+    if shutdown_ray:
+        import ray  # Lazy import Ray
+
+        ray.shutdown()
+    gc.collect()
+
+    if not current_omni_platform.is_cpu():
+        current_omni_platform.empty_cache()
+        try:
+            torch._C._host_emptyCache()
+        except AttributeError:
+            logger.warning("torch._C._host_emptyCache() only available in Pytorch >=2.5")
+
+
 def _split_request_config_by_per_output_sizes(cfg: dict[str, Any]) -> list[dict[str, Any]] | None:
     """If ``extra_body`` has list ``height``/``width``, return one config per index (scalar h/w, ``num_outputs_per_prompt=1``)."""
     eb = cfg.get("extra_body")
@@ -88,13 +125,6 @@ PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None
 PromptImageInput = list[Any] | Any | None
 PromptVideoInput = list[Any] | Any | None
 
-try:
-    from vllm.distributed.parallel_state import cleanup_dist_env_and_memory  # type: ignore
-except Exception:  # pragma: no cover
-
-    def cleanup_dist_env_and_memory() -> None:
-        return None
-
 
 def get_open_port(host: str = "127.0.0.1", *, max_attempts: int = 128) -> int:
     """Return a local TCP port that is suitable for binding a new listener.