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

[VLM] Support VLM ViT Piecewise CUDA Graph (#14422)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
Co-authored-by: kousakawang <wanghanpei@bytedance.com>
This commit is contained in:
Yuan Luo
2025-12-13 20:49:40 +08:00
committed by GitHub
parent f1bbd26ff7
commit 3b8a824b8b
4 changed files with 690 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -51,7 +51,7 @@ from sglang.srt.layers.linear import (
from sglang.srt.layers.quantization import QuantizationConfig
from sglang.srt.layers.rotary_embedding import apply_rotary_pos_emb
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import add_prefix
from sglang.srt.utils import add_prefix, get_bool_env_var
ROTARY_EMBED_CLASSES = {
"normal": apply_rotary_pos_emb,
@@ -273,6 +273,10 @@ class VisionTritonAttention(nn.Module):
**kwargs,
):
super().__init__()
use_data_parallel = (
kwargs["use_data_parallel"] if "use_data_parallel" in kwargs else False
)
self.tp_size = 1 if use_data_parallel else get_attention_tp_size()
def forward(
self,
@@ -290,22 +294,42 @@ class VisionTritonAttention(nn.Module):
Returns:
[b * s, h, head_size]
"""
cu_seqlens = resolve_seqlens(cu_seqlens, bsz, seq_len, device=q.device)
if get_bool_env_var("SGLANG_VIT_ENABLE_CUDA_GRAPH") and self.tp_size == 1:
if "output_ws" not in kwargs:
raise RuntimeError("output_ws should be prepared for cuda-graph mode")
# [b * s, head, head_size]
output = torch.empty_like(q)
seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
max_seqlen = seq_lens.max().item()
context_attention_fwd(
q,
k,
v,
output,
cu_seqlens.cuda(),
seq_lens.cuda(),
max_seqlen,
is_causal=False,
)
if not isinstance(cu_seqlens, list):
raise RuntimeError("cuda-graph mode cu_seqlens should be a list")
output = kwargs["output_ws"]
context_attention_fwd(
q,
k,
v,
output,
cu_seqlens[0],
cu_seqlens[1],
cu_seqlens[2],
is_causal=False,
)
else:
cu_seqlens = resolve_seqlens(cu_seqlens, bsz, seq_len, device=q.device)
# [b * s, head, head_size]
output = torch.empty_like(q)
seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
max_seqlen = seq_lens.max().item()
context_attention_fwd(
q,
k,
v,
output,
cu_seqlens.cuda(),
seq_lens.cuda(),
max_seqlen,
is_causal=False,
)
return output
@@ -318,6 +342,10 @@ class VisionFlash3Attention(nn.Module):
if not _is_cuda:
raise Exception("VisionFlash3Attention is only available for cuda")
super().__init__()
use_data_parallel = (
kwargs["use_data_parallel"] if "use_data_parallel" in kwargs else False
)
self.tp_size = 1 if use_data_parallel else get_attention_tp_size()
def forward(
self,
@@ -335,21 +363,32 @@ class VisionFlash3Attention(nn.Module):
Returns:
[b * s, h, head_size]
"""
cu_seqlens = resolve_seqlens(cu_seqlens, bsz, seq_len, device=q.device)
if get_bool_env_var("SGLANG_VIT_ENABLE_CUDA_GRAPH") and self.tp_size == 1:
max_seqlen = cu_seqlens[1]
output = flash_attn_varlen_func(
q,
k,
v,
cu_seqlens_q=cu_seqlens[0],
cu_seqlens_k=cu_seqlens[0],
max_seqlen_q=max_seqlen,
max_seqlen_k=max_seqlen,
)
else:
cu_seqlens = resolve_seqlens(cu_seqlens, bsz, seq_len, device=q.device)
cu_seqlens = cu_seqlens.to(dtype=torch.int32).to(q.device)
seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
max_seqlen = seq_lens.max().item()
cu_seqlens = cu_seqlens.to(dtype=torch.int32).to(q.device)
seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
max_seqlen = seq_lens.max().item()
output = flash_attn_varlen_func(
q,
k,
v,
cu_seqlens_q=cu_seqlens,
cu_seqlens_k=cu_seqlens,
max_seqlen_q=max_seqlen,
max_seqlen_k=max_seqlen,
)
output = flash_attn_varlen_func(
q,
k,
v,
cu_seqlens_q=cu_seqlens,
cu_seqlens_k=cu_seqlens,
max_seqlen_q=max_seqlen,
max_seqlen_k=max_seqlen,
)
return output
@@ -545,6 +584,7 @@ class VisionAttention(nn.Module):
dropout=dropout,
flatten_batch=flatten_batch,
softmax_in_single_precision=softmax_in_single_precision,
use_data_parallel=use_data_parallel,
)
self.use_qkv_parallel = use_qkv_parallel
@@ -652,6 +692,8 @@ class VisionAttention(nn.Module):
bsz, s, _ = x_shape
head = self.num_attention_heads_per_partition
kv_head = self.num_attention_kv_heads_per_partition
attn_output_ws = kwargs["output_ws"] if "output_ws" in kwargs else None
if self.use_qkv_parallel:
# [b, s, embed_dim] --> [b, s, embed_dim]
qkv, _ = self.qkv_proj(x)
@@ -729,6 +771,7 @@ class VisionAttention(nn.Module):
seq_len=s,
cu_seqlens=cu_seqlens,
attention_mask=attention_mask,
output_ws=attn_output_ws,
)
assert output.dim() == 3, output.shape

84
python/sglang/srt/models/qwen2_5_vl.py
View File

@@ -72,8 +72,9 @@ from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.qwen2 import Qwen2Model
from sglang.srt.models.utils import RotaryPosMixin, permute_inv
from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model
from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import add_prefix
from sglang.srt.utils import add_prefix, get_bool_env_var
logger = logging.getLogger(__name__)
@@ -167,7 +168,9 @@ class Qwen2_5_VisionBlock(nn.Module):
x: torch.Tensor,
cu_seqlens: torch.Tensor,
position_embeddings: torch.Tensor,
output_ws=None,
) -> torch.Tensor:
ws = output_ws
S, B, H = x.shape
# norm1: flatten to 2D -> [S*B, H], then reshape back
x2d = x.reshape(-1, H)
@@ -179,6 +182,7 @@ class Qwen2_5_VisionBlock(nn.Module):
hidden_states,
cu_seqlens=cu_seqlens,
position_embeddings=position_embeddings,
output_ws=ws,
)
attn = rearrange(attn, "b s h -> s b h")
@@ -256,6 +260,7 @@ class Qwen2_5_VisionTransformer(nn.Module, RotaryPosMixin):
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
use_data_parallel: bool = False,
max_context_len: Optional[int] = None,
) -> None:
super().__init__()
@@ -308,6 +313,13 @@ class Qwen2_5_VisionTransformer(nn.Module, RotaryPosMixin):
use_data_parallel=use_data_parallel,
)
# Resource prepared for vit cuda graph
self.tp_size = (
1 if use_data_parallel else get_tensor_model_parallel_world_size()
)
self.max_context_len = max_context_len
self.cuda_graph_runner: Optional[ViTCudaGraphRunner] = ViTCudaGraphRunner(self)
def get_window_index(self, grid_thw):
cu_window_seqlens: list = [0]
window_index_id = 0
@@ -378,6 +390,9 @@ class Qwen2_5_VisionTransformer(nn.Module, RotaryPosMixin):
x: torch.Tensor,
grid_thw: torch.Tensor,
) -> torch.Tensor:
if get_bool_env_var("SGLANG_VIT_ENABLE_CUDA_GRAPH") and self.tp_size == 1:
return self.forward_with_cuda_graph(x, grid_thw)
# patchify
x = x.to(device=self.device, dtype=self.dtype)
x = self.patch_embed(x)
@@ -446,6 +461,72 @@ class Qwen2_5_VisionTransformer(nn.Module, RotaryPosMixin):
return x
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)
# compute position embedding
rotary_pos_emb = self.rot_pos_emb(grid_thw)
window_index, cu_window_seqlens = self.get_window_index(grid_thw)
cu_window_seqlens = torch.tensor(
cu_window_seqlens,
device=x.device,
dtype=torch.int32,
)
cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens)
window_index = window_index.to(device=x.device)
reverse_indices = permute_inv(window_index)
rotary_pos_emb = rotary_pos_emb.to(device=x.device, dtype=x.dtype)
# patch token num
seq_len, _ = x.size()
# [G, M, hidden]
x = x.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
x = x[window_index, :, :] # [G, M, hidden]
x = x.reshape(seq_len, -1) # [seq_len, hidden]
rotary_pos_emb = rotary_pos_emb.reshape(
seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1
)
rotary_pos_emb = rotary_pos_emb[window_index, :, :]
rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
position_embeddings = (emb.cos(), emb.sin())
# After building position_embeddings, make sure both cos and sin are on
# the same device/dtype as the attention input
position_embeddings = (
position_embeddings[0].to(x.device, x.dtype),
position_embeddings[1].to(x.device, x.dtype),
)
# compute cu_seqlens - move cu_seqlens to GPU and make it int32
cu_seqlens = torch.cat(
[
torch.tensor([0], device=x.device, dtype=torch.int32),
(grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2])
.cumsum(dim=0)
.to(device=x.device, dtype=torch.int32),
]
)
cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens])
return self.cuda_graph_runner.run(
x=x,
position_embeddings=position_embeddings,
cu_seqlens=cu_seqlens,
cu_window_seqlens=cu_window_seqlens,
output_indices=reverse_indices,
)
class Qwen2_5_VLForConditionalGeneration(nn.Module):
# BitandBytes specific attributes
@@ -485,6 +566,7 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module):
quant_config=quant_config,
prefix=add_prefix("visual", prefix),
use_data_parallel=self.use_data_parallel,
max_context_len=self.config.max_position_embeddings,
)
self.model = Qwen2Model(

263
python/sglang/srt/multimodal/vit_cuda_graph_runner.py Normal file
View File

@@ -0,0 +1,263 @@
# 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 CUDA Graph Runner class."""
from __future__ import annotations
import inspect
from typing import Dict, Hashable, Optional, Tuple
import torch
import torch.nn as nn
from sglang.srt.layers.attention.vision import VisionAttention
from sglang.srt.server_args import get_global_server_args
class ViTCudaGraphRunner:
"""ViT CUDA Graph Runner
Cached with graph_key = seq_len, for each seq_len capture once.
expose run(), internally call create_graph().
exceed call invokes replay().
"""
def __init__(
self,
vit: nn.Module,
) -> None:
self.vit = vit
# graph_key -> buffers / graphs
self.block_input: Dict[Hashable, torch.Tensor] = {}
self.block_ws: Dict[Hashable, torch.Tensor] = {}
self.block_graphs: Dict[Hashable, torch.cuda.CUDAGraph] = {}
self.block_output: Dict[Hashable, torch.Tensor] = {}
# captured seqlens buffers (addresses must be stable for cuda-graph replay)
self.cu_full_len: Dict[Hashable, torch.Tensor] = {}
self.cu_window_len: Dict[Hashable, torch.Tensor] = {}
self.cu_full_len_kk: Dict[Hashable, torch.Tensor] = {}
self.cu_window_len_kk: Dict[Hashable, torch.Tensor] = {}
# rotary position buffers shared across graphs
self.sin_cos_ws: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
self.max_context_len = getattr(vit, "max_context_len", None)
self._fullatt_block_indexes = set(getattr(vit, "fullatt_block_indexes", ()))
first_blk = vit.blocks[0]
self._blk_accepts_output_ws = (
"output_ws" in inspect.signature(first_blk.forward).parameters
)
self._attn: Optional[VisionAttention] = getattr(first_blk, "attn", None)
self._attn_backend = getattr(self._attn, "qkv_backend", None)
@property
def device(self) -> torch.device:
return self.vit.device
@property
def dtype(self) -> torch.dtype:
return self.vit.dtype
def _ensure_sin_cos_ws(self, seq_len: int, head_dim: int):
if self.sin_cos_ws is None:
max_shape = self.max_context_len or seq_len
max_shape = max(max_shape, seq_len)
cos_ws = torch.empty(
max_shape, head_dim, dtype=self.dtype, device=self.device
)
sin_ws = torch.empty(
max_shape, head_dim, dtype=self.dtype, device=self.device
)
self.sin_cos_ws = (cos_ws, sin_ws)
else:
if self.sin_cos_ws[0].size(0) < seq_len:
max_shape = max(self.sin_cos_ws[0].size(0) * 2, seq_len)
cos_ws = torch.empty(
max_shape, head_dim, dtype=self.dtype, device=self.device
)
sin_ws = torch.empty(
max_shape, head_dim, dtype=self.dtype, device=self.device
)
self.sin_cos_ws = (cos_ws, sin_ws)
def _get_graph_key(self, x_3d: torch.Tensor) -> int:
# x_3d: [S, B, H], B=1, S as graph_key
return x_3d.shape[0]
def _create_graph(self, graph_key: int, temp_cos_sin):
graph = torch.cuda.CUDAGraph()
vit = self.vit
cu_window = self.cu_window_len[graph_key]
cu_full = self.cu_full_len[graph_key]
cu_window_kk = self.cu_window_len_kk[graph_key]
cu_full_kk = self.cu_full_len_kk[graph_key]
max_full_len = int(cu_full_kk.max().item())
max_window_len = int(cu_window_kk.max().item())
override_backend = get_global_server_args().mm_attention_backend
with torch.cuda.graph(graph):
y = None
for layer_num, blk in enumerate(vit.blocks):
if layer_num in vit.fullatt_block_indexes:
cu_seqlens_now = cu_full
cu_seqlens_kk_now = cu_full_kk
max_len = max_full_len
else:
cu_seqlens_now = cu_window
cu_seqlens_kk_now = cu_window_kk
max_len = max_window_len
if override_backend == "triton_attn":
cu_seq_len_ws = [cu_seqlens_now, cu_seqlens_kk_now, max_len]
elif override_backend == "fa3":
cu_seq_len_ws = [cu_seqlens_now, max_len]
else:
raise RuntimeError("Not supported ViT attention backend")
if layer_num == 0:
y = blk(
self.block_input[graph_key],
cu_seqlens=cu_seq_len_ws,
position_embeddings=temp_cos_sin,
output_ws=self.block_ws[graph_key],
)
else:
y = blk(
y,
cu_seqlens=cu_seq_len_ws,
position_embeddings=temp_cos_sin,
output_ws=self.block_ws[graph_key],
)
self.block_output[graph_key] = vit.merger(y)
self.block_graphs[graph_key] = graph
def create_graph(
self,
x_3d: torch.Tensor, # [S, 1, H]
position_embeddings: Tuple[torch.Tensor, torch.Tensor], # (cos, sin), [S, D]
cu_seqlens: torch.Tensor,
cu_window_seqlens: torch.Tensor,
) -> int:
vit = self.vit
graph_key = self._get_graph_key(x_3d)
if graph_key in self.block_graphs:
return graph_key
# make sure rotary workspace
head_dim = position_embeddings[0].shape[1]
self._ensure_sin_cos_ws(graph_key, head_dim)
used_cos_ws = self.sin_cos_ws[0][:graph_key, :]
used_sin_ws = self.sin_cos_ws[1][:graph_key, :]
used_cos_ws.copy_(position_embeddings[0])
used_sin_ws.copy_(position_embeddings[1])
temp_cos_sin = (used_cos_ws, used_sin_ws)
# 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] = torch.empty_like(
x_3d, device=self.device
).contiguous()
self.block_input[graph_key] = torch.empty_like(
x_3d, device=self.device
).contiguous()
self.block_ws[graph_key] = torch.empty(
graph_key,
num_heads,
attn_head_dim,
device=self.device,
dtype=self.dtype,
)
if graph_key not in self.cu_window_len:
self.cu_window_len[graph_key] = cu_window_seqlens
self.cu_full_len[graph_key] = cu_seqlens
self.cu_window_len_kk[graph_key] = (
cu_window_seqlens[1:] - cu_window_seqlens[:-1]
)
self.cu_full_len_kk[graph_key] = cu_seqlens[1:] - cu_seqlens[:-1]
self._create_graph(graph_key, temp_cos_sin)
return graph_key
def replay(
self,
graph_key: int,
x_3d: torch.Tensor,
position_embeddings: Tuple[torch.Tensor, torch.Tensor],
output_indices: Optional[torch.Tensor] = None,
) -> torch.Tensor:
# update rotary workspace content
head_dim = position_embeddings[0].shape[1]
self._ensure_sin_cos_ws(graph_key, head_dim)
used_cos_ws = self.sin_cos_ws[0][:graph_key, :]
used_sin_ws = self.sin_cos_ws[1][:graph_key, :]
used_cos_ws.copy_(position_embeddings[0])
used_sin_ws.copy_(position_embeddings[1])
# 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,
position_embeddings: Tuple[torch.Tensor, torch.Tensor],
cu_seqlens: torch.Tensor,
cu_window_seqlens: torch.Tensor,
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,
position_embeddings=position_embeddings,
cu_seqlens=cu_seqlens,
cu_window_seqlens=cu_window_seqlens,
)
return self.replay(
graph_key=graph_key,
x_3d=x_3d,
position_embeddings=position_embeddings,
output_indices=output_indices,
)