[NPU] [Diffusion] Fix critical Ascend NPU Diffusion regression/bugs & restore 2-NPU CI testcase (#34855)

Co-authored-by: Elizaveta Martirosian <elizabet3000@mail.ru>
Co-authored-by: Arseniy Mironov <98156294+Napkin-AI@users.noreply.github.com>
Co-authored-by: Alexandr <117110413+Allor-maker@users.noreply.github.com>
Co-authored-by: P_Alex_Tr <aleksandr.smyshlaev@yandex.ru>

.github/workflows/pr-test-npu.yml

@@ -167,6 +167,10 @@ jobs:
     needs: [check-changes, pr-gate, set-image-config]
     if: ${{ !failure() && !cancelled() && needs.check-changes.outputs.multimodal_gen == 'true' }}
     runs-on: linux-aarch64-a3-800t-2
+    strategy:
+      fail-fast: false
+      matrix:
+        part: [0, 1]
     container:
       image: ${{ needs.set-image-config.outputs.CANN_image_a3 }}
     steps:
@@ -209,21 +213,28 @@ jobs:
           SGLANG_DIFFUSION_ARTIFACT_DIR: ${{ github.workspace }}/diffusion-failures
         run: |
           cd python
-          python3 sglang/multimodal_gen/test/run_suite.py --suite 1-npu
+          python3 sglang/multimodal_gen/test/run_suite.py \
+            --suite 1-npu \
+            --partition-id ${{ matrix.part }} \
+            --total-partitions 2
 
       - name: Upload diffusion failure artifacts
         if: always()
         uses: actions/upload-artifact@v4
         with:
-          name: diffusion-failures-npu-1-${{ github.run_attempt }}
+          name: diffusion-failures-npu-1-part${{ matrix.part }}-${{ github.run_attempt }}
           path: diffusion-failures/
           if-no-files-found: ignore
           retention-days: 7
 
   multimodal-gen-test-2-npu-a3:
     needs: [check-changes, pr-gate, set-image-config]
-    if: false  # Temporarily disabled; will be restored once the issue is fixed.
-    runs-on: linux-aarch64-a3-800t-2
+    if: ${{ !failure() && !cancelled() && needs.check-changes.outputs.multimodal_gen == 'true' }}
+    runs-on: linux-aarch64-a3-800t-16
+    strategy:
+      fail-fast: false
+      matrix:
+        part: [0, 1]
     container:
       image: ${{ needs.set-image-config.outputs.CANN_image_a3 }}
     steps:
@@ -266,13 +277,16 @@ jobs:
           SGLANG_DIFFUSION_ARTIFACT_DIR: ${{ github.workspace }}/diffusion-failures
         run: |
           cd python
-          python3 sglang/multimodal_gen/test/run_suite.py --suite 2-npu
+          python3 sglang/multimodal_gen/test/run_suite.py \
+            --suite 2-npu \
+            --partition-id ${{ matrix.part }} \
+            --total-partitions 2
 
       - name: Upload diffusion failure artifacts
         if: always()
         uses: actions/upload-artifact@v4
         with:
-          name: diffusion-failures-npu-2-${{ github.run_attempt }}
+          name: diffusion-failures-npu-2-part${{ matrix.part }}-${{ github.run_attempt }}
           path: diffusion-failures/
           if-no-files-found: ignore
           retention-days: 7