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

.github/workflows/diffusion-ci-gt-gen-npu.yml

@@ -88,7 +88,7 @@ jobs:
     if: |
       needs.compute-diffusion-partitions.result == 'success' &&
       needs.compute-diffusion-partitions.outputs.matrix-1npu != '{"include":[]}'
-    runs-on: linux-aarch64-a3-2
+    runs-on: linux-aarch64-a3-800t-2
     strategy:
       fail-fast: false
       matrix: ${{ fromJson(needs.compute-diffusion-partitions.outputs.matrix-1npu) }}
@@ -98,6 +98,8 @@ jobs:
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.ref || github.ref }}
 
       - name: Mark repository safe
         run: |
@@ -158,7 +160,7 @@ jobs:
     if: |
       needs.compute-diffusion-partitions.result == 'success' &&
       needs.compute-diffusion-partitions.outputs.matrix-2npu != '{"include":[]}'
-    runs-on: linux-aarch64-a3-16
+    runs-on: linux-aarch64-a3-800t-16
     strategy:
       fail-fast: false
       matrix: ${{ fromJson(needs.compute-diffusion-partitions.outputs.matrix-2npu) }}
@@ -168,6 +170,8 @@ jobs:
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.ref || github.ref }}
 
       - name: Mark repository safe
         run: |

.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

.github/workflows/release-docker-npu-nightly.yml

@@ -80,6 +80,6 @@ jobs:
           push: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }}
           provenance: false
           build-args: |
-            SGLANG_KERNEL_NPU_TAG=2026.8.10
+            SGLANG_KERNEL_NPU_TAG=2026.8.13
             CANN_VERSION=${{ matrix.cann_version }}
             DEVICE_TYPE=${{ matrix.device_type }}

.github/workflows/release-docker-npu.yml

@@ -87,7 +87,7 @@ jobs:
           push: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }}
           provenance: false
           build-args: |
-            SGLANG_KERNEL_NPU_TAG=2026.8.10
+            SGLANG_KERNEL_NPU_TAG=2026.8.13
             CANN_VERSION=${{ matrix.cann_version }}
             DEVICE_TYPE=${{ matrix.device_type }}
             SGLANG_TAG=${{ steps.version.outputs.version }}

docker/npu.Dockerfile

@@ -115,7 +115,7 @@ RUN ${PIP_INSTALL} wheel==0.45.1 pybind11 pyyaml decorator scipy attrs psutil \
     && cd sgl-kernel-npu \
     && wget https://github.com/sgl-project/sgl-kernel-npu/releases/download/${SGLANG_KERNEL_NPU_TAG}/sgl-kernel-npu-${SGLANG_KERNEL_NPU_TAG}-torch2.10.0-py311-cann${CANN_VERSION}-${DEVICE_TYPE}-$(arch).zip \
     && unzip sgl-kernel-npu-${SGLANG_KERNEL_NPU_TAG}-torch2.10.0-py311-cann${CANN_VERSION}-${DEVICE_TYPE}-$(arch).zip \
-    && ${PIP_INSTALL} deep_ep*.whl sgl_kernel_npu*.whl \
+    && ${PIP_INSTALL} deep_ep*.whl sgl_kernel_npu*.whl attentions*.whl \
     && cd .. && rm -rf sgl-kernel-npu \
     && cd "$(python3 -m pip show deep-ep | awk '/^Location:/ {print $2}')" && ln -sf deep_ep/deep_ep_cpp*.so
 

python/sglang/kernels/ops/attention/flash_attention.py

@@ -257,7 +257,6 @@ def flash_attn_varlen_func(
     ver=3,
     out=None,
 ):
-
     if ver == 3:
         return fa3_flash_attn_varlen_func(
             q,

python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py

@@ -1,3 +1,4 @@
+from collections.abc import Sequence
 from dataclasses import dataclass
 from typing import Any
 
@@ -15,6 +16,143 @@
 logger = init_logger(__name__)
 
 
+def _packed_boundaries(
+    cu_seqlens: torch.Tensor,
+    cu_seqlens_host: Sequence[int] | None,
+    total_tokens: int,
+    name: str,
+) -> tuple[int, ...]:
+    if cu_seqlens is None:
+        raise ValueError(f"{name} is required for NPU packed attention")
+    if cu_seqlens.ndim != 1 or cu_seqlens.dtype not in (
+        torch.int32,
+        torch.int64,
+    ):
+        raise ValueError(f"{name} must be a 1D int32 or int64 tensor")
+    if cu_seqlens_host is not None and len(cu_seqlens_host) != cu_seqlens.numel():
+        raise ValueError(f"{name} and its host copy must have the same length")
+
+    boundaries = tuple(
+        int(value)
+        for value in (
+            cu_seqlens.tolist() if cu_seqlens_host is None else cu_seqlens_host
+        )
+    )
+    if len(boundaries) < 2 or boundaries[0] != 0:
+        raise ValueError(f"{name} must start with 0 and contain at least one sequence")
+    if boundaries[-1] != total_tokens:
+        raise ValueError(
+            f"{name} must end at the packed token count {total_tokens}, "
+            f"got {boundaries[-1]}"
+        )
+    if any(stop < start for start, stop in zip(boundaries[:-1], boundaries[1:])):
+        raise ValueError(f"{name} must be non-decreasing")
+    return boundaries
+
+
+def fused_infer_attention_varlen(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    *,
+    cu_seqlens_q_host: Sequence[int] | None = None,
+    cu_seqlens_k_host: Sequence[int] | None = None,
+    softmax_scale: float | None = None,
+    return_softmax_lse: bool = False,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    tensors = {"q": q, "k": k, "v": v}
+    invalid_layouts = [name for name, tensor in tensors.items() if tensor.ndim != 3]
+    if invalid_layouts:
+        raise ValueError(
+            "NPU packed attention requires q, k, and v in [T, N, D] layout; "
+            f"invalid tensors: {', '.join(invalid_layouts)}"
+        )
+    invalid_devices = [
+        name
+        for name, tensor in tensors.items()
+        if tensor.device.type != "npu" or tensor.device != q.device
+    ]
+    if invalid_devices:
+        raise ValueError(
+            "NPU packed attention requires q, k, and v on the same NPU; "
+            f"invalid tensors: {', '.join(invalid_devices)}"
+        )
+    if not (q.dtype == k.dtype == v.dtype):
+        raise ValueError(
+            "NPU packed attention requires q, k, and v with the same dtype"
+        )
+    if k.shape[:2] != v.shape[:2]:
+        raise ValueError(
+            "NPU packed attention requires matching K/V token and head counts"
+        )
+    if q.shape[-1] != k.shape[-1]:
+        raise ValueError("NPU packed attention requires matching Q/K head dimensions")
+
+    q_boundaries = _packed_boundaries(
+        cu_seqlens_q, cu_seqlens_q_host, q.shape[0], "cu_seqlens_q"
+    )
+    k_boundaries = _packed_boundaries(
+        cu_seqlens_k, cu_seqlens_k_host, k.shape[0], "cu_seqlens_k"
+    )
+    if len(q_boundaries) != len(k_boundaries):
+        raise ValueError("cu_seqlens_q and cu_seqlens_k must describe the same batch")
+
+    q_nonempty = [
+        stop > start for start, stop in zip(q_boundaries[:-1], q_boundaries[1:])
+    ]
+    k_nonempty = [
+        stop > start for start, stop in zip(k_boundaries[:-1], k_boundaries[1:])
+    ]
+    if q_nonempty != k_nonempty:
+        raise NotImplementedError(
+            "NPU packed attention does not support a sequence that is empty only "
+            "on the query or key/value side"
+        )
+    actual_seq_lengths = [
+        stop for stop, nonempty in zip(q_boundaries[1:], q_nonempty) if nonempty
+    ]
+    actual_seq_lengths_kv = [
+        stop for stop, nonempty in zip(k_boundaries[1:], k_nonempty) if nonempty
+    ]
+    if not actual_seq_lengths:
+        output = torch.empty_like(q)
+        if return_softmax_lse:
+            lse = torch.empty(
+                (q.shape[1], q.shape[0]), dtype=torch.float32, device=q.device
+            )
+            return output, lse
+        return output
+
+    if not (q.is_contiguous() and k.is_contiguous() and v.is_contiguous()):
+        if q.shape == k.shape == v.shape:
+            q, k, v = torch.stack((q, k, v), dim=0).unbind(0)
+        else:
+            q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
+
+    output, lse = torch.ops.npu.npu_fused_infer_attention_score(
+        q,
+        k,
+        v,
+        num_heads=q.shape[1],
+        num_key_value_heads=k.shape[1],
+        scale=q.shape[-1] ** -0.5 if softmax_scale is None else softmax_scale,
+        input_layout="TND",
+        actual_seq_lengths=actual_seq_lengths,
+        actual_seq_lengths_kv=actual_seq_lengths_kv,
+        softmax_lse_flag=return_softmax_lse,
+    )
+    if not return_softmax_lse:
+        return output
+    if lse.shape != (q.shape[0], q.shape[1], 1):
+        raise RuntimeError(
+            "Unexpected Ascend TND softmax LSE shape: "
+            f"expected {(q.shape[0], q.shape[1], 1)}, got {tuple(lse.shape)}"
+        )
+    return output, lse.squeeze(-1).transpose(0, 1).contiguous()
+
+
 @dataclass
 class AscendFAMetadata:
     pass
@@ -106,5 +244,71 @@ def forward(
         )
         output = output.transpose(1, 2)
         if return_softmax_lse:
-            return output, lse
+            return output, lse.squeeze(-1)
         return output
+
+    def forward_varlen(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        *,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: int,
+        cu_seqlens_host: tuple[int, ...] | None = None,
+    ) -> torch.Tensor:
+        del max_seqlen
+        if self.causal:
+            bounds = (
+                cu_seqlens_host
+                if cu_seqlens_host is not None
+                else tuple(int(item) for item in cu_seqlens.tolist())
+            )
+            output = torch.empty_like(query)
+            for start, stop in zip(bounds[:-1], bounds[1:]):
+                if start == stop:
+                    continue
+                segment = self.forward(
+                    query[start:stop].unsqueeze(0),
+                    key[start:stop].unsqueeze(0),
+                    value[start:stop].unsqueeze(0),
+                    None,
+                )
+                output[start:stop].copy_(segment[0])
+            return output
+
+        return fused_infer_attention_varlen(
+            query,
+            key,
+            value,
+            cu_seqlens,
+            cu_seqlens,
+            cu_seqlens_q_host=cu_seqlens_host,
+            cu_seqlens_k_host=cu_seqlens_host,
+            softmax_scale=self.softmax_scale,
+        )
+
+    def forward_ring_kv_chunk(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Run one Ascend TND ring chunk and return LSE in ``[H, Tq]``."""
+        cu_seqlens_q = torch.tensor(
+            [0, query.shape[0]], dtype=torch.int32, device=query.device
+        )
+        cu_seqlens_k = torch.tensor(
+            [0, key.shape[0]], dtype=torch.int32, device=key.device
+        )
+        return fused_infer_attention_varlen(
+            query,
+            key,
+            value,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            cu_seqlens_q_host=(0, query.shape[0]),
+            cu_seqlens_k_host=(0, key.shape[0]),
+            softmax_scale=self.softmax_scale,
+            return_softmax_lse=True,
+        )

python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py

@@ -209,6 +209,21 @@ def forward_varlen(
             f"{type(self).__name__} does not implement packed varlen attention"
         )
 
+    def forward_ring_kv_chunk(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Attend local queries to one rotated KV chunk for ring merging.
+
+        Inputs use packed ``[T, H, D]`` layout. The returned attention output
+        has the query shape and softmax LSE uses ``[H, Tq]`` layout.
+        """
+        raise NotImplementedError(
+            f"{type(self).__name__} does not implement ring KV-chunk attention"
+        )
+
 
 def wrap_attention_impl_forward(attn_impl: AttentionImpl) -> AttentionImpl:
     return wrap_method_with_debug_kernel_once(