[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/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.
+    if: ${{ !failure() && !cancelled() && needs.check-changes.outputs.multimodal_gen == 'true' }}
-    runs-on: linux-aarch64-a3-800t-2
+    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 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 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 @@ class AscendFAImpl(AttentionImpl):
         )
         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 @@ class AttentionImpl(ABC, Generic[T]):
             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(

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

@@ -472,3 +472,37 @@ class FlashAttentionImpl(AttentionImpl):
             ver=fa_ver,
         )
         return output[0] if isinstance(output, tuple) else output
+
+    def forward_ring_kv_chunk(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Run one non-causal FlashAttention ring chunk with LSE output."""
+        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
+        )
+        result = flash_attn_varlen_func(
+            query,
+            key,
+            value,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=query.shape[0],
+            max_seqlen_k=key.shape[0],
+            softmax_scale=self.softmax_scale,
+            causal=False,
+            ver=fa_ver,
+            return_softmax_lse=True,
+        )
+        if not isinstance(result, tuple):
+            raise RuntimeError(
+                "FlashAttention did not return the softmax LSE required by ring "
+                "attention"
+            )
+        output, softmax_lse, *_ = result
+        return output, softmax_lse

python/sglang/multimodal_gen/runtime/layers/attention/layer.py

@@ -1303,7 +1303,7 @@ class USPAttention(nn.Module):
             q.squeeze(0),
             k.squeeze(0),
             v.squeeze(0),
-            softmax_scale=self.softmax_scale,
+            attn_impl=self.attn_impl,
             real_seq_len=int(attn_mask_meta["pad_start"]),
             ring_ws=get_ring_parallel_world_size(),
         )

python/sglang/multimodal_gen/runtime/layers/usp.py

@@ -8,7 +8,6 @@ import torch.distributed as dist
 import torch.distributed._functional_collectives as ft_c
 from torch.distributed.tensor.experimental._attention import _cp_options
 
-from sglang.kernels.ops.attention.flash_attention import flash_attn_varlen_func
 from sglang.kernels.ops.diffusion import pack_qkv_destination_major, usp_merge_heads
 from sglang.multimodal_gen.runtime.distributed.parallel_state import (
     get_ring_ctx,
@@ -16,9 +15,6 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
     get_ulysses_parallel_rank,
     get_ulysses_parallel_world_size,
 )
-from sglang.multimodal_gen.runtime.layers.attention.backends import (
-    flash_attn as _fa_backend,
-)
 from sglang.srt.utils.common import torch_release
 
 _cp_options.enable_load_balance = False
@@ -821,7 +817,7 @@ def _ring_attention_varlen(
     k: torch.Tensor,
     v: torch.Tensor,
     *,
-    softmax_scale: float,
+    attn_impl: "AttentionImpl",
     real_seq_len: int,
     ring_ws: int,
 ) -> torch.Tensor:
@@ -860,7 +856,6 @@ def _ring_attention_varlen(
     kv_bufs = [kv0, torch.empty_like(kv0)]
     cur = 0
 
-    q_cu = torch.tensor([0, ring_chunk_len], dtype=torch.int32, device=q.device)
     out_acc: torch.Tensor | None = None
     lse_acc: torch.Tensor | None = None
     pending_ops = None
@@ -891,27 +886,11 @@ def _ring_attention_varlen(
             max(real_seq_len - src_rank * ring_chunk_len, 0), ring_chunk_len
         )
         if remote_used > 0:
-            k_cu = torch.tensor([0, remote_used], dtype=torch.int32, device=q.device)
+            step_out, step_lse = attn_impl.forward_ring_kv_chunk(
-            result = flash_attn_varlen_func(
                 q,
                 kv_bufs[cur][0, :remote_used],
                 kv_bufs[cur][1, :remote_used],
-                cu_seqlens_q=q_cu,
-                cu_seqlens_k=k_cu,
-                max_seqlen_q=ring_chunk_len,
-                max_seqlen_k=remote_used,
-                softmax_scale=softmax_scale,
-                causal=False,
-                ver=_fa_backend.fa_ver,
-                return_softmax_lse=True,
             )
-            if not isinstance(result, tuple):
-                raise RuntimeError(
-                    "flash_attn_varlen_func did not return softmax_lse; ring "
-                    "parallelism requires a backend that supports "
-                    "return_softmax_lse=True."
-                )
-            step_out, step_lse, *_ = result
             out_acc, lse_acc = _ring_merge_attention(
                 out_acc, lse_acc, step_out, step_lse
             )

python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py

@@ -239,12 +239,16 @@ def module_weight_bytes(module) -> int:
     seen: set[int] = set()
     total = 0
     for tensor in list(module.parameters()) + list(module.buffers()):
-        storage = tensor.untyped_storage()
+        try:
-        pointer = storage.data_ptr()
+            storage = tensor.untyped_storage()
+            pointer = storage.data_ptr()
+            storage_bytes = storage.nbytes()
+        except RuntimeError:
+            continue
         if pointer == 0 or pointer in seen:
             continue
         seen.add(pointer)
-        total += storage.nbytes()
+        total += storage_bytes
     return total
 
 

python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py

@@ -552,7 +552,7 @@ def _minimax_h3_attention_core_impl(
             q,
             k,
             v,
-            softmax_scale=attention.softmax_scale,
+            attn_impl=attention._attention_impl,
             real_seq_len=max_seqlen,
             ring_ws=ring_ws,
         )

python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py

@@ -179,7 +179,7 @@ class SelfAttention(nn.Module):
         # USPAttention handles SP communication internally; the tail meta keeps
         # SP padding out of the softmax.
         out = self.attn(q, k, v, attn_mask_meta=attn_mask_meta)
-        out = out.view(b, s, -1)
+        out = out.reshape(b, s, -1)
 
         out, _ = self.o(out)
         return out

python/sglang/multimodal_gen/test/server/ascend/perf_baselines_npu.json

@@ -73,7 +73,7 @@
         "flux_2_image_t2i_2npu": {
             "stages_ms": {
                 "InputValidationStage": 0.08,
-                "TextEncodingStage": 192.4,
+                "TextEncodingStage": 363.09,
                 "ImageVAEEncodingStage": 0.01,
                 "LatentPreparationStage": 0.97,
                 "TimestepPreparationStage": 34.65,
@@ -134,7 +134,8 @@
             },
             "expected_e2e_ms": 46557.7,
             "expected_avg_denoise_ms": 872.7,
-            "expected_median_denoise_ms": 905.81
+            "expected_median_denoise_ms": 905.81,
+            "estimated_full_test_time_s": 488.4
         },
         "wan2_1_t2v_1.3b_1_npu": {
             "stages_ms": {
@@ -257,7 +258,8 @@
             },
             "expected_e2e_ms": 193947.19,
             "expected_avg_denoise_ms": 4691.07,
-            "expected_median_denoise_ms": 4773.22
+            "expected_median_denoise_ms": 4773.22,
+            "estimated_full_test_time_s": 987.8
         },
         "qwen_image_t2i_2npu": {
             "stages_ms": {
@@ -322,14 +324,15 @@
             },
             "expected_e2e_ms": 34362.34,
             "expected_avg_denoise_ms": 610.41,
-            "expected_median_denoise_ms": 615.39
+            "expected_median_denoise_ms": 615.39,
+            "estimated_full_test_time_s": 275.2
         },
         "ernie_image_t2i_1npu": {
             "stages_ms": {
                 "InputValidationStage": 0.07,
                 "PromptEnhancementStage": 8787.2,
                 "TextEncodingStage": 35.87,
-                "TimestepPreparationStage": 32.74,
+                "TimestepPreparationStage": 246.30,
                 "LatentPreparationStage": 0.2,
                 "DenoisingStage": 47471.93,
                 "DecodingStage": 42.24
@@ -572,7 +575,8 @@
             },
             "expected_e2e_ms": 109909.47,
             "expected_avg_denoise_ms": 2636.6,
-            "expected_median_denoise_ms": 2705.05
+            "expected_median_denoise_ms": 2705.05,
+            "estimated_full_test_time_s": 239.7
          },
        "mova_360p_ti2va_2npu": {
            "stages_ms": {
@@ -645,13 +649,13 @@
             "stages_ms": {
                 "InputValidationStage": 3.8,
                 "TextEncodingStage": 1029.19,
-                "LTX2TextConnectorStage": 56.3,
+                "LTX2TextConnectorStage":  654.62,
                 "LTX2SigmaPreparationStage": 0.19,
                 "TimestepPreparationStage": 33.52,
                 "LTX2AVLatentPreparationStage": 0.45,
                 "LTX2ImageEncodingStage": 93.32,
                 "LTX2AVDenoisingStage": 29672.06,
-                "LTX2AVDecodingStage": 560.9
+                "LTX2AVDecodingStage": 1240.88
             },
             "denoise_step_ms": {
                 "0": 544.81,
@@ -697,7 +701,8 @@
             },
             "expected_e2e_ms": 31484.37,
             "expected_avg_denoise_ms": 741.58,
-            "expected_median_denoise_ms": 746.28
+            "expected_median_denoise_ms": 746.28,
+            "estimated_full_test_time_s": 255.2
         }
     }
 }

python/sglang/multimodal_gen/test/test_utils.py

@@ -45,7 +45,7 @@ SGL_TEST_FILES_CI_DATA_REVISION = "15b30030ef980756788ab40072f9223fe21a5526"
 # The NPU pin is kept as a separate branch so ascend GT can be bumped independently
 # when it's regenerated on its own cadence.
 if current_platform.is_npu():
-    SGL_TEST_FILES_CI_DATA_REVISION = "d180ad38872dff3d1ad03e4610cffcda874d3eb8"
+    SGL_TEST_FILES_CI_DATA_REVISION = "8e3d717e65fb87339c2974382a092a731669f884"
 
 SGL_TEST_FILES_CONSISTENCY_GT_ROOT = (
     "https://raw.githubusercontent.com/"
@@ -1632,20 +1632,29 @@ def _remote_file_exists(url: str) -> bool | None:
 
 def _load_remote_gt_image(url: str) -> np.ndarray:
     last_error: Exception | None = None
-    for _ in range(3):
+    attempts = 3
+    for attempt in range(1, attempts + 1):
         try:
             resp = requests.get(url, timeout=60)
             try:
                 if resp.status_code == 200:
-                    image = Image.open(io.BytesIO(resp.content)).convert("RGB")
+                    with Image.open(io.BytesIO(resp.content)) as image:
-                    return np.array(image)
+                        return np.array(image.convert("RGB"))
                 last_error = FileNotFoundError(f"GT image not found: {url}")
                 if resp.status_code not in (403, 429) and resp.status_code < 500:
                     break
             finally:
                 resp.close()
-        except requests.RequestException as exc:
+        except (OSError, ValueError, requests.RequestException) as exc:
             last_error = exc
+        if attempt < attempts:
+            logger.warning(
+                "GT image download failed (attempt %d/%d), retrying: %s",
+                attempt,
+                attempts,
+                url,
+            )
+            time.sleep(attempt)
     raise FileNotFoundError(f"GT image not found: {url}") from last_error
 
 

python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py

@@ -239,6 +239,26 @@ class TestModuleWeightBytes:
         module.register_buffer("b", backing[512:])
         assert module_weight_bytes(module) == 4096
 
+    def test_invalid_storage_is_skipped(self):
+        """Offloaded models may expose a tensor whose storage cannot be queried."""
+
+        class InvalidStorage:
+            def data_ptr(self):
+                raise RuntimeError("invalid python storage")
+
+        class InvalidStorageTensor:
+            def untyped_storage(self):
+                return InvalidStorage()
+
+        class ModuleWithInvalidStorage:
+            def parameters(self):
+                return iter((torch.ones(4), InvalidStorageTensor()))
+
+            def buffers(self):
+                return iter(())
+
+        assert module_weight_bytes(ModuleWithInvalidStorage()) == 4 * 4
+
 
 class TestPinBenefit:
     def test_a_stepped_component_counts_every_step(self):

python/sglang/multimodal_gen/test/unit/test_usp_ring_tail_pad.py

@@ -22,6 +22,7 @@ class TestRingTailPadDispatch(unittest.TestCase):
         obj.backend = AttentionBackendEnum.FA
         obj.causal = False
         obj.dropout_p = 0.0
+        obj.attn_impl = object()
         return obj
 
     def test_tail_pad_meta_reaches_the_ring_kernel(self):
@@ -30,12 +31,12 @@ class TestRingTailPadDispatch(unittest.TestCase):
         meta = {"pad_start": 13, "pad_end": 16, "local_pad": 3}
         seen = {}
 
-        def fake_ring(qc, kc, vc, *, softmax_scale, real_seq_len, ring_ws):
+        def fake_ring(qc, kc, vc, *, attn_impl, real_seq_len, ring_ws):
             seen.update(
                 shape=tuple(qc.shape),
                 real=real_seq_len,
                 ws=ring_ws,
-                scale=softmax_scale,
+                impl=attn_impl,
             )
             return torch.ones_like(qc)
 
@@ -61,6 +62,7 @@ class TestRingTailPadDispatch(unittest.TestCase):
         self.assertEqual(out.shape, q.shape)
         self.assertEqual(seen["real"], 13)
         self.assertEqual(seen["ws"], 2)
+        self.assertIs(seen["impl"], obj.attn_impl)
         self.assertEqual(seen["shape"], (4, 2, 8))
         # Last ring rank holds global rows [12, 16): row 13 onward is pad.
         self.assertTrue(torch.all(out[0, 1:] == 0))

python/sglang/srt/hardware_backend/npu/utils.py

@@ -101,12 +101,15 @@ def init_npu_backend():
         logger.warning("NPU custom kernel packages unavailable: %s", e)
 
     import torch_npu
-    from torch_npu.contrib import transfer_to_npu  # noqa: F401
 
-    # Re-mock torch.cuda.is_available cuz transfer_to_npu mocks it True
+    # These imports lead to unpredictable behavior in diffusion models
-    torch.cuda.is_available = lambda: False
+    # and a significant reduction in performance.
+    if "sglang.multimodal_gen" not in sys.modules:
+        from torch_npu.contrib import transfer_to_npu  # noqa: F401
 
-    torch_npu.npu.config.allow_internal_format = True
+        # Re-mock torch.cuda.is_available cuz transfer_to_npu mocks it True
+        torch.cuda.is_available = lambda: False
+        torch_npu.npu.config.allow_internal_format = True
     torch_npu.npu.set_compile_mode(jit_compile=False)
 
 

scripts/ci/npu/npu_ci_install_dependency.sh

@@ -60,11 +60,11 @@ ${PIP_INSTALL} triton-ascend==3.2.1.dev20260530 --extra-index-url=https://mirror
 
 
 ### Install sgl-kernel-npu
-SGLANG_KERNEL_NPU_TAG="2026.8.10"
+SGLANG_KERNEL_NPU_TAG="2026.8.13"
 mkdir sgl-kernel-npu
 (cd sgl-kernel-npu && wget "${GITHUB_PROXY_URL:=""}https://github.com/sgl-project/sgl-kernel-npu/releases/download/${SGLANG_KERNEL_NPU_TAG}/sgl-kernel-npu-${SGLANG_KERNEL_NPU_TAG}-torch${PYTORCH_VERSION}-py311-cann9.0.0-${DEVICE_TYPE}-$(arch).zip" \
 && unzip ./sgl-kernel-npu-${SGLANG_KERNEL_NPU_TAG}-torch${PYTORCH_VERSION}-py311-cann9.0.0-${DEVICE_TYPE}-$(arch).zip \
-&& ${UV_PIP_INSTALL} ./deep_ep*.whl ./sgl_kernel_npu*.whl \
+&& ${UV_PIP_INSTALL} ./deep_ep*.whl ./sgl_kernel_npu*.whl ./attentions*.whl \
 && (cd "$(python3 -m pip show deep-ep | grep -E '^Location:' | awk '{print $2}')" && ln -s deep_ep/deep_ep_cpp*.so))
 
 ### Install custom-ops