From b98d472158f707c11b39510fa8406253bbfd5008 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=90=D1=80=D1=82=D0=B5=D0=BC=20=D0=A1=D0=B0=D0=B2=D0=BA?= =?UTF-8?q?=D0=B8=D0=BD?= <58187114+OrangeRedeng@users.noreply.github.com> Date: Sat, 22 Aug 2026 18:32:20 +0300 Subject: [PATCH] [NPU] [Diffusion] Fix critical Ascend NPU Diffusion regression/bugs & restore 2-NPU CI testcase (#34855) Co-authored-by: Elizaveta Martirosian 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 --- .github/workflows/diffusion-ci-gt-gen-npu.yml | 8 +- .github/workflows/pr-test-npu.yml | 26 ++- .../workflows/release-docker-npu-nightly.yml | 2 +- .github/workflows/release-docker-npu.yml | 2 +- docker/npu.Dockerfile | 2 +- .../kernels/ops/attention/flash_attention.py | 1 - .../layers/attention/backends/ascend_fa.py | 206 +++++++++++++++++- .../attention/backends/attention_backend.py | 15 ++ .../layers/attention/backends/flash_attn.py | 34 +++ .../runtime/layers/attention/layer.py | 2 +- .../multimodal_gen/runtime/layers/usp.py | 25 +-- .../memory_managers/host_memory_budget.py | 10 +- .../runtime/models/dits/minimax_h3.py | 2 +- .../runtime/models/dits/mova_video_dit.py | 2 +- .../server/ascend/perf_baselines_npu.json | 23 +- .../sglang/multimodal_gen/test/test_utils.py | 19 +- .../test/unit/test_host_memory_budget.py | 20 ++ .../test/unit/test_usp_ring_tail_pad.py | 6 +- .../sglang/srt/hardware_backend/npu/utils.py | 11 +- scripts/ci/npu/npu_ci_install_dependency.sh | 4 +- 20 files changed, 356 insertions(+), 64 deletions(-) diff --git a/.github/workflows/diffusion-ci-gt-gen-npu.yml b/.github/workflows/diffusion-ci-gt-gen-npu.yml index 3dd4b323a..36e569714 100644 --- a/.github/workflows/diffusion-ci-gt-gen-npu.yml +++ b/.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: | diff --git a/.github/workflows/pr-test-npu.yml b/.github/workflows/pr-test-npu.yml index 214b8a332..70d6f5b2b 100644 --- a/.github/workflows/pr-test-npu.yml +++ b/.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 diff --git a/.github/workflows/release-docker-npu-nightly.yml b/.github/workflows/release-docker-npu-nightly.yml index 1904f2ec6..02a551e23 100644 --- a/.github/workflows/release-docker-npu-nightly.yml +++ b/.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 }} diff --git a/.github/workflows/release-docker-npu.yml b/.github/workflows/release-docker-npu.yml index 87f4d0488..56012eca2 100644 --- a/.github/workflows/release-docker-npu.yml +++ b/.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 }} diff --git a/docker/npu.Dockerfile b/docker/npu.Dockerfile index e8922e23f..a2ba523e5 100644 --- a/docker/npu.Dockerfile +++ b/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 diff --git a/python/sglang/kernels/ops/attention/flash_attention.py b/python/sglang/kernels/ops/attention/flash_attention.py index 5e9e39837..2a146ce88 100644 --- a/python/sglang/kernels/ops/attention/flash_attention.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py index aebf21f91..b25388ae4 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py +++ b/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, + ) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py index 41a080258..733b3bbe4 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py +++ b/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( diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py index d42e1936b..5710e321d 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index b721a5899..c25b8ff90 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/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(), ) diff --git a/python/sglang/multimodal_gen/runtime/layers/usp.py b/python/sglang/multimodal_gen/runtime/layers/usp.py index b98550edb..766cc23e3 100644 --- a/python/sglang/multimodal_gen/runtime/layers/usp.py +++ b/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) - result = flash_attn_varlen_func( + step_out, step_lse = attn_impl.forward_ring_kv_chunk( 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 ) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py index 82d5f4fa5..ca986e801 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py +++ b/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() - pointer = storage.data_ptr() + try: + 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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py index b84a0655a..fcd66de55 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/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, ) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py b/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py index cc13c42a2..4a0b4de0c 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py +++ b/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 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 1708db6bb..0b2d604bc 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 @@ -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 } } } diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index d5b42f310..79e91233b 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/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") - return np.array(image) + with Image.open(io.BytesIO(resp.content)) as 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 diff --git a/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py b/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py index 646c738a1..5caeb91e3 100644 --- a/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py +++ b/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): diff --git a/python/sglang/multimodal_gen/test/unit/test_usp_ring_tail_pad.py b/python/sglang/multimodal_gen/test/unit/test_usp_ring_tail_pad.py index a8be07e6a..7079677e9 100644 --- a/python/sglang/multimodal_gen/test/unit/test_usp_ring_tail_pad.py +++ b/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)) diff --git a/python/sglang/srt/hardware_backend/npu/utils.py b/python/sglang/srt/hardware_backend/npu/utils.py index 664bb2881..b4a6968f4 100644 --- a/python/sglang/srt/hardware_backend/npu/utils.py +++ b/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 - torch.cuda.is_available = lambda: False + # These imports lead to unpredictable behavior in diffusion models + # 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) diff --git a/scripts/ci/npu/npu_ci_install_dependency.sh b/scripts/ci/npu/npu_ci_install_dependency.sh index 83023face..10b9c1570 100755 --- a/scripts/ci/npu/npu_ci_install_dependency.sh +++ b/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