From fc543df2899582729ef23cbcd9f9362393624222 Mon Sep 17 00:00:00 2001 From: cen121212 Date: Mon, 9 Mar 2026 10:13:35 +0800 Subject: [PATCH] [NPU] qwen3_vl encoder support graph --- .../npu/graph_runner/vit_npu_graph_runner.py | 229 ++++++++++++++++++ python/sglang/srt/layers/attention/vision.py | 23 +- python/sglang/srt/models/qwen3_vl.py | 93 +++++-- 3 files changed, 313 insertions(+), 32 deletions(-) create mode 100644 python/sglang/srt/hardware_backend/npu/graph_runner/vit_npu_graph_runner.py diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/vit_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/vit_npu_graph_runner.py new file mode 100644 index 000000000..87cf281a9 --- /dev/null +++ b/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, + ) diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 4abe5b2c4..77a8cde46 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/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 diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index 34a645078..38a4767b0 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/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 = 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) + ( + x, + cu_seqlens, + rotary_pos_emb_cos, + rotary_pos_emb_sin, + ) = self._prepare_graph_inputs(x, 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.cuda_graph_runner.run( + return self.graph_runners.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)