[NPU] qwen3_vl encoder support graph

python/sglang/srt/hardware_backend/npu/graph_runner/vit_npu_graph_runner.py

@@ -0,0 +1,229 @@
+# Copyright 2023-2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""ViT NPU Graph Runner class."""
+
+from __future__ import annotations
+
+from typing import Dict, Hashable, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch_npu
+
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    set_graph_pool_id,
+)
+from sglang.srt.layers.attention.vision import VisionAttention
+from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner
+from sglang.srt.server_args import get_global_server_args
+
+
+class ViTNpuGraphRunner(ViTCudaGraphRunner):
+    """Generic ViT NPU Graph Runner.
+
+    This runner captures the "blocks + merger + deepstack merger (optional)" part
+    of a vision transformer into a NPU graph and replays it for identical shapes.
+
+    Optional for Qwen3 deepstack:
+      - vit.deepstack_vision_indexes: Sequence[int]
+      - vit.deepstack_merger_list: nn.ModuleList (same length as deepstack_vision_indexes)
+    """
+
+    _graph_memory_pool = None
+
+    def __init__(
+        self,
+        vit: nn.Module,
+    ) -> None:
+        super().__init__(vit)
+        self.device_module = torch.get_device_module(self.device)
+        self.cu_seq_lens: Dict[Hashable, torch.Tensor] = {}
+
+        # rotary position buffers shared across graphs
+        self.sin_cos_ws: Dict[Hashable, Tuple[torch.Tensor, torch.Tensor]] = {}
+
+    @property
+    def device(self) -> torch.device:
+        return self.vit.device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.vit.dtype
+
+    def _create_graph(
+        self,
+        graph_key: int,
+    ):
+
+        graph = torch_npu.npu.NPUGraph()
+        vit = self.vit
+
+        override_backend = get_global_server_args().mm_attention_backend
+        with torch_npu.npu.graph(graph, pool=ViTNpuGraphRunner._graph_memory_pool):
+            y = None
+            deepstack_outs: List[torch.Tensor] = []
+            deepstack_capture_idx = 0
+
+            for layer_num, blk in enumerate(vit.blocks):
+                if override_backend == "ascend_attn":
+                    cu_seq_lens = self.cu_seq_lens[graph_key]
+                else:
+                    raise RuntimeError("Not supported ViT attention backend")
+
+                if layer_num == 0:
+                    y = blk(
+                        self.block_input[graph_key],
+                        cu_seqlens=cu_seq_lens,
+                        rotary_pos_emb_cos=self.sin_cos_ws[graph_key][0],
+                        rotary_pos_emb_sin=self.sin_cos_ws[graph_key][1],
+                        output_ws=self.block_ws[graph_key],
+                    )
+                else:
+                    y = blk(
+                        y,
+                        cu_seqlens=cu_seq_lens,
+                        rotary_pos_emb_cos=self.sin_cos_ws[graph_key][0],
+                        rotary_pos_emb_sin=self.sin_cos_ws[graph_key][1],
+                        output_ws=self.block_ws[graph_key],
+                    )
+
+                # Optional deepstack support (Qwen3-VL)
+                if (
+                    self._deepstack_visual_indexes
+                    and layer_num in self._deepstack_visual_indexes
+                ):
+                    if self._deepstack_merger_list is None:
+                        raise RuntimeError(
+                            "deepstack_visual_indexes exists but deepstack_merger_list is missing."
+                        )
+                    deepstack_out = self._deepstack_merger_list[deepstack_capture_idx](
+                        y
+                    )
+                    deepstack_outs.append(deepstack_out)
+                    deepstack_capture_idx += 1
+
+            main_out = vit.merger(y)
+
+            if deepstack_outs:
+                self.block_output[graph_key] = torch.cat(
+                    [main_out] + deepstack_outs, dim=1
+                )
+            else:
+                self.block_output[graph_key] = main_out
+
+        self.block_graphs[graph_key] = graph
+
+    def create_graph(
+        self,
+        x_3d: torch.Tensor,  # [S, 1, H]
+        cu_seqlens: torch.Tensor,
+        rotary_pos_emb_cos: Optional[torch.Tensor] = None,
+        rotary_pos_emb_sin: Optional[torch.Tensor] = None,
+    ) -> int:
+        vit = self.vit
+        graph_key = self._get_graph_key(x_3d)
+
+        if graph_key in self.block_graphs:
+            return graph_key
+
+        if ViTNpuGraphRunner._graph_memory_pool is None:
+            ViTNpuGraphRunner._graph_memory_pool = (
+                self.device_module.graph_pool_handle()
+            )
+        # Set graph pool id globally to be able to use symmetric memory
+        set_graph_pool_id(ViTNpuGraphRunner._graph_memory_pool)
+
+        # pre-allocate workspace
+        attn_module: VisionAttention = vit.blocks[0].attn
+        num_heads = attn_module.num_attention_heads_per_partition
+        attn_head_dim = attn_module.head_size
+
+        if graph_key not in self.block_output:
+            self.block_output[graph_key] = x_3d
+            self.block_input[graph_key] = x_3d
+            self.block_ws[graph_key] = torch.empty(
+                graph_key,
+                num_heads,
+                attn_head_dim,
+                device=self.device,
+                dtype=self.dtype,
+            )
+            if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None:
+                self.sin_cos_ws[graph_key] = (rotary_pos_emb_cos, rotary_pos_emb_sin)
+
+        if graph_key not in self.cu_seq_lens:
+            seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+            self.cu_seq_lens[graph_key] = seq_lens.to("cpu").to(torch.int32)
+
+        if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None:
+            self._create_graph(
+                graph_key=graph_key,
+            )
+
+        return graph_key
+
+    def replay(
+        self,
+        graph_key: int,
+        x_3d: torch.Tensor,
+        rotary_pos_emb_cos: Optional[torch.Tensor] = None,
+        rotary_pos_emb_sin: Optional[torch.Tensor] = None,
+        output_indices: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None:
+            # update rotary workspace content
+            self.sin_cos_ws[graph_key][0].copy_(rotary_pos_emb_cos)
+            self.sin_cos_ws[graph_key][1].copy_(rotary_pos_emb_sin)
+
+        # copy input
+        self.block_input[graph_key].copy_(x_3d)
+
+        # replay
+        self.block_graphs[graph_key].replay()
+
+        out = self.block_output[graph_key]
+
+        # Optional output reordering (Qwen2.5-VL window permutation inverse)
+        if output_indices is not None:
+            out = out.index_select(0, output_indices)
+
+        return out
+
+    def run(
+        self,
+        x: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        rotary_pos_emb_cos: Optional[torch.Tensor] = None,
+        rotary_pos_emb_sin: Optional[torch.Tensor] = None,
+        output_indices: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        # x: [seq_len, hidden] -> [S, B=1, H]
+        x_3d = x.unsqueeze(1)
+        graph_key = self._get_graph_key(x_3d)
+        if graph_key not in self.block_graphs:
+            self.create_graph(
+                x_3d=x_3d,
+                cu_seqlens=cu_seqlens,
+                rotary_pos_emb_cos=rotary_pos_emb_cos,
+                rotary_pos_emb_sin=rotary_pos_emb_sin,
+            )
+
+        return self.replay(
+            graph_key=graph_key,
+            x_3d=x_3d,
+            rotary_pos_emb_cos=rotary_pos_emb_cos,
+            rotary_pos_emb_sin=rotary_pos_emb_sin,
+            output_indices=output_indices,
+        )

python/sglang/srt/layers/attention/vision.py

@@ -678,28 +678,33 @@ class VisionAscendAttention(nn.Module):
         Returns:
              [b * s, h, head_size]
         """
-        cu_seqlens = resolve_seqlens(cu_seqlens, bsz, seq_len, device="cpu")
+        if envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get():
+            if "output_ws" not in kwargs:
+                raise RuntimeError("output_ws should be prepared for npu-graph mode")
+            output = kwargs["output_ws"]
+            # graph mode: runner already passes seq_lens (int32 on CPU)
+            seq_len_arg = cu_seqlens
+        else:
+            cu_seqlens = resolve_seqlens(cu_seqlens, bsz, seq_len, device="cpu")
+            seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+            if seq_lens.is_npu:
+                seq_lens = seq_lens.to("cpu")
+            output = torch.empty_like(q)
+            seq_len_arg = seq_lens.to(torch.int32)
 
-        seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
-        if seq_lens.is_npu:
-            # cu_seqlens must be on cpu because of operator restriction
-            seq_lens = seq_lens.to("cpu")
         _, num_heads, head_size = q.shape
         num_kv_heads = k.shape[1]
-        output = torch.empty_like(q)
 
-        # operator requires pta version >= 2.5.1
         torch_npu._npu_flash_attention_unpad(
             query=q,
             key=k,
             value=v,
-            seq_len=seq_lens.to(torch.int32),
+            seq_len=seq_len_arg,
             scale_value=head_size**-0.5,
             num_heads=num_heads,
             num_kv_heads=num_kv_heads,
             out=output,
         )
-
         return output
 
 

python/sglang/srt/models/qwen3_vl.py

@@ -17,6 +17,7 @@
 import logging
 import math
 import re
+from collections import defaultdict
 from functools import lru_cache, partial
 from typing import Callable, Iterable, List, Optional, Tuple, Union
 
@@ -74,6 +75,16 @@ from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import add_prefix, get_int_env_var, is_npu, round_up
 from sglang.srt.utils.hf_transformers_utils import get_processor
 
+_is_npu = is_npu()
+graph_runners_dict = defaultdict(lambda: ViTCudaGraphRunner)
+if _is_npu:
+    from sglang.srt.hardware_backend.npu.graph_runner.vit_npu_graph_runner import (
+        ViTNpuGraphRunner,
+    )
+
+    graph_runners_dict["npu"] = ViTNpuGraphRunner
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -335,7 +346,7 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
 
         workspace_buffer = None
         if get_global_server_args().mm_attention_backend == "flashinfer_cudnn":
-            if torch.cuda.is_available() and (not is_npu()):
+            if torch.cuda.is_available() and (not _is_npu):
                 ws_device = torch.device("cuda", torch.cuda.current_device())
             else:
                 ws_device = self.device
@@ -390,7 +401,7 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
         self.tp_size = (
             1 if use_data_parallel else get_tensor_model_parallel_world_size()
         )
-        self.cuda_graph_runner: Optional[ViTCudaGraphRunner] = ViTCudaGraphRunner(self)
+        self.graph_runners = graph_runners_dict[self.device.type](self)
 
     @property
     def dtype(self) -> torch.dtype:
@@ -721,6 +732,8 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
         grid_thw: torch.Tensor,
     ) -> torch.Tensor:
         if envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get():
+            if _is_npu:
+                return self.forward_with_npu_graph(x, grid_thw)
             return self.forward_with_cuda_graph(x, grid_thw)
 
         x = x.to(device=self.device, dtype=self.dtype)
@@ -788,7 +801,7 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
         else:
             sequence_lengths = None
             cu_seqlens = torch.from_numpy(token_cu_seqlens)
-            if not is_npu():
+            if not _is_npu:
                 cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)
             else:
                 cu_seqlens = cu_seqlens.to("cpu")
@@ -823,37 +836,45 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
         )  # [seq_len, hidden_size * (1 + depth_of_deepstack)]
         return hidden_states
 
+    def forward_with_npu_graph(
+        self,
+        x: torch.Tensor,
+        grid_thw: torch.Tensor,
+    ) -> torch.Tensor:
+        (
+            x,
+            cu_seqlens,
+            rotary_pos_emb_cos,
+            rotary_pos_emb_sin,
+        ) = self._prepare_graph_inputs(x, grid_thw)
+
+        cu_seqlens = cu_seqlens.to("cpu")
+        return self.graph_runners.run(
+            x=x,
+            rotary_pos_emb_cos=rotary_pos_emb_cos,
+            rotary_pos_emb_sin=rotary_pos_emb_sin,
+            cu_seqlens=cu_seqlens,
+            output_indices=None,
+        )
+
     def forward_with_cuda_graph(
         self,
         x: torch.Tensor,
         grid_thw: torch.Tensor,
     ) -> torch.Tensor:
-        # patchify
+        (
-        x = x.to(device=self.device, dtype=self.dtype)
+            x,
-        x = self.patch_embed(x)
+            cu_seqlens,
-
+            rotary_pos_emb_cos,
-        if isinstance(grid_thw, list):
+            rotary_pos_emb_sin,
-            grid_thw_list = grid_thw
+        ) = self._prepare_graph_inputs(x, grid_thw)
-            grid_thw = torch.tensor(grid_thw, dtype=torch.int32)
-        else:
-            grid_thw_list = grid_thw.tolist()
-
-        pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
-        x += pos_embeds
-
-        # rotary embedding -> (cos, sin)
-        rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw_list)
-
-        # compute cu_seqlens
-        cu_seqlens = compute_cu_seqlens_from_grid_numpy(grid_thw)
         if not isinstance(cu_seqlens, torch.Tensor):
             cu_seqlens = torch.tensor(cu_seqlens, device=x.device, dtype=torch.int32)
         else:
             cu_seqlens = cu_seqlens.to(device=x.device, dtype=torch.int32)
         cu_seqlens = cu_seqlens.contiguous()
 
-        # blocks + merger + deepstack(optional) via CUDA Graph Runner
+        return self.graph_runners.run(
-        return self.cuda_graph_runner.run(
             x=x,
             position_embeddings=None,
             rotary_pos_emb_cos=rotary_pos_emb_cos,
@@ -890,6 +911,32 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
             loaded_params.add(name)
         return loaded_params
 
+    def _prepare_graph_inputs(self, x: torch.Tensor, grid_thw: torch.Tensor) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+    ]:
+        # patchify
+        x = x.to(device=self.device, dtype=self.dtype)
+        x = self.patch_embed(x)
+
+        if isinstance(grid_thw, list):
+            grid_thw_list = grid_thw
+            grid_thw = torch.tensor(grid_thw, dtype=torch.int32)
+        else:
+            grid_thw_list = grid_thw.tolist()
+
+        pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
+        x += pos_embeds
+
+        # rotary embedding -> (cos, sin)
+        rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw_list)
+
+        # compute cu_seqlens
+        cu_seqlens = compute_cu_seqlens_from_grid_numpy(grid_thw)
+        return x, cu_seqlens, rotary_pos_emb_cos, rotary_pos_emb_sin
+
 
 cached_get_processor = lru_cache(get_processor)