chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[NPU] qwen3_vl encoder support graph

Browse Source
This commit is contained in:
cen121212
2026-03-09 10:13:35 +08:00
committed by GitHub
parent 8c5ca37aef
commit fc543df289
3 changed files with 313 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

229
python/sglang/srt/hardware_backend/npu/graph_runner/vit_npu_graph_runner.py Normal file
View File

@@ -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,
)

23
python/sglang/srt/layers/attention/vision.py
View File

@@ -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

93
python/sglang/srt/models/qwen3_vl.py
View File

@@ -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)