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

[VLM] Introduce FlashInfer CUDNN Prefill as ViT Backend (#19003)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-02-24 19:49:22 +08:00
committed by GitHub
parent 1a83b2c15d
commit 31c7dc9d99
4 changed files with 678 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -34,6 +34,7 @@ _is_npu = is_npu()
_is_hip = is_hip()
if _is_cuda:
from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
from sgl_kernel.flash_attn import flash_attn_varlen_func
if _is_npu:
@@ -64,6 +65,24 @@ ROTARY_EMBED_CLASSES = {
"normal": apply_rotary_pos_emb,
}
# === Vision Encoder === #
FLASHINFER_WORKSPACE_SIZE_BYTES = 128 * 1024 * 1024
# Batch buckets for cuDNN graph caching - graphs are cached per bucket size
# This avoids creating a new graph for each unique batch size at runtime
BATCH_BUCKETS = [8, 16, 32, 64]
# Bucketized max seqlens to reduce cuDNN recompilation frequency while
# preserving a tighter upper bound than a single fixed max seqlen.
FLASHINFER_MAX_SEQLEN_BUCKETS = [
4 * 1024,
8 * 1024,
16 * 1024,
32 * 1024,
64 * 1024,
128 * 1024,
]
@dataclasses.dataclass
class SingletonCache:
@@ -452,6 +471,128 @@ class VisionFlash4Attention(nn.Module):
return output
class VisionFlashInferAttention(nn.Module):
def __init__(
self,
**kwargs,
):
if not _is_cuda:
raise Exception("VisionFlashInferAttention is only available for cuda")
super().__init__()
self.workspace_buffer = (
kwargs["workspace_buffer"] if "workspace_buffer" in kwargs else None
)
def forward(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
cu_seqlens: torch.Tensor | SingletonCache | None,
bsz: int,
seq_len: int,
**kwargs,
) -> torch.Tensor:
r"""
Args:
cu_seqlens: [b]
Returns:
[b * s, h, head_size]
"""
if "sequence_lengths" not in kwargs:
raise RuntimeError(
"sequence_lengths should be prepared for vision flashinfer_cudnn attention backend"
)
if "max_seqlen" not in kwargs:
raise RuntimeError(
"max_seqlen should be prepared for vision flashinfer_cudnn attention backend"
)
sequence_lengths = kwargs["sequence_lengths"] # (B_padded,) or (B_padded,1,1,1)
max_seqlen = kwargs["max_seqlen"]
# max_seqlen must be python int
if isinstance(max_seqlen, torch.Tensor):
if max_seqlen.is_cuda:
max_seqlen = int(max_seqlen.detach().cpu().item())
else:
max_seqlen = int(max_seqlen.item())
else:
max_seqlen = int(max_seqlen)
# flatten if caller gives (b, s, h, d)
is_reshaped = q.dim() == 4
if is_reshaped:
reshape_batch_size = q.shape[0]
q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
if not isinstance(cu_seqlens, torch.Tensor):
raise RuntimeError(
"flashinfer_cudnn expects packed indptrs as a torch.Tensor"
)
# sequence_lengths -> (B,)
if not isinstance(sequence_lengths, torch.Tensor):
raise RuntimeError("sequence_lengths must be a torch.Tensor")
seq_lens_1d = sequence_lengths.view(-1).to(device=q.device, dtype=torch.int32)
B = int(seq_lens_1d.numel())
# cu_seqlens contains packed *element indptrs*:
# [qk_indptr(B+1), v_indptr(B+1), o_indptr(B+1)] => total 3*(B+1)
cu_seqlens_1d = cu_seqlens.view(-1).to(device=q.device, dtype=torch.int32)
expected = 3 * (B + 1)
if int(cu_seqlens_1d.numel()) != expected:
raise RuntimeError(
f"packed indptr numel mismatch: got {cu_seqlens_1d.numel()}, expected {expected} (= 3*(B+1))"
)
split = B + 1
indptr_qk = cu_seqlens_1d[:split].view(split, 1, 1, 1)
indptr_v = cu_seqlens_1d[split : 2 * split].view(split, 1, 1, 1)
indptr_o = cu_seqlens_1d[2 * split :].view(split, 1, 1, 1)
# cuDNN style: (B,1,1,1)
seq_lens_4d = seq_lens_1d.view(B, 1, 1, 1)
# indptr are in ELEMENT offsets (not token offsets)
token_width_q = int(q.shape[1] * q.shape[2]) # heads * head_dim on this rank
total_elems_q = int(q.numel())
# check each real sequence fits
# (skip padded tail where seq_len==0)
start_elems = indptr_qk.view(-1)[:-1] # (B,)
end_elems = start_elems + seq_lens_1d * token_width_q
if (end_elems > total_elems_q).any():
raise RuntimeError("offset + len out of bounds; packed indptr is wrong")
_, _, head_size = q.shape
scale = head_size**-0.5
output, _ = cudnn_batch_prefill_with_kv_cache(
q,
k,
v,
scale,
self.workspace_buffer,
max_token_per_sequence=max_seqlen,
max_sequence_kv=max_seqlen,
actual_seq_lens_q=seq_lens_4d,
actual_seq_lens_kv=seq_lens_4d,
causal=False,
return_lse=True,
batch_offsets_q=indptr_qk,
batch_offsets_k=indptr_qk,
batch_offsets_v=indptr_v,
batch_offsets_o=indptr_o,
is_cuda_graph_compatible=True,
)
if is_reshaped:
output = rearrange(output, "(b s) h d -> b s h d", b=reshape_batch_size)
return output
class VisionAiterAttention(nn.Module):
def __init__(
self,
@@ -552,6 +693,7 @@ QKV_BACKEND_IMPL = {
"sdpa": VisionSdpaAttention,
"fa3": VisionFlash3Attention,
"fa4": VisionFlash4Attention,
"flashinfer_cudnn": VisionFlashInferAttention,
"ascend_attn": VisionAscendAttention,
"aiter_attn": VisionAiterAttention,
}
@@ -594,6 +736,7 @@ class VisionAttention(nn.Module):
use_data_parallel: bool = False,
use_dp_attention_reduce: bool = False,
aux_stream: Optional[torch.cuda.Stream] = None,
workspace_buffer: Optional[torch.Tensor] = None,
**kwargs,
):
super().__init__()
@@ -651,6 +794,7 @@ class VisionAttention(nn.Module):
flatten_batch=flatten_batch,
softmax_in_single_precision=softmax_in_single_precision,
use_data_parallel=use_data_parallel,
workspace_buffer=workspace_buffer,
)
self.use_qkv_parallel = use_qkv_parallel
@@ -686,6 +830,8 @@ class VisionAttention(nn.Module):
prefix=add_prefix("proj", prefix),
use_dp_attention_reduce=use_dp_attention_reduce,
)
self.workspace_buffer = workspace_buffer
self.aux_stream = aux_stream
self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] if aux_stream else []
@@ -829,6 +975,10 @@ class VisionAttention(nn.Module):
kv_head = self.num_attention_kv_heads_per_partition
attn_output_ws = kwargs["output_ws"] if "output_ws" in kwargs else None
max_seqlen = kwargs["max_seqlen"] if "max_seqlen" in kwargs else None
sequence_lengths = (
kwargs["sequence_lengths"] if "sequence_lengths" in kwargs else None
)
if self.use_qkv_parallel:
# [b, s, embed_dim] --> [b, s, embed_dim]
qkv, _ = self.qkv_proj(x)
@@ -935,6 +1085,8 @@ class VisionAttention(nn.Module):
seq_len=s,
cu_seqlens=cu_seqlens,
attention_mask=attention_mask,
sequence_lengths=sequence_lengths,
max_seqlen=max_seqlen,
output_ws=attn_output_ws,
)

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

@@ -30,7 +30,12 @@ from sglang.srt.configs.qwen3_vl import Qwen3VLConfig, Qwen3VLVisionConfig
from sglang.srt.distributed import get_tensor_model_parallel_world_size
from sglang.srt.distributed.parallel_state import get_pp_group
from sglang.srt.environ import envs
from sglang.srt.layers.attention.vision import VisionAttention
from sglang.srt.layers.attention.vision import (
BATCH_BUCKETS,
FLASHINFER_MAX_SEQLEN_BUCKETS,
FLASHINFER_WORKSPACE_SIZE_BYTES,
VisionAttention,
)
from sglang.srt.layers.dp_attention import (
get_attention_tp_rank,
get_attention_tp_size,
@@ -66,15 +71,12 @@ from sglang.srt.models.utils import (
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, get_int_env_var, is_npu
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
logger = logging.getLogger(__name__)
# === Vision Encoder === #
class Qwen3_VisionMLP(nn.Module):
def __init__(
@@ -161,6 +163,7 @@ class Qwen3_VisionBlock(nn.Module):
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
use_data_parallel: bool = False,
workspace_buffer: torch.Tensor | None = None,
) -> None:
super().__init__()
if norm_layer is None:
@@ -179,6 +182,7 @@ class Qwen3_VisionBlock(nn.Module):
prefix=add_prefix("attn", prefix),
use_data_parallel=use_data_parallel,
use_dp_attention_reduce=is_dp_attention_enabled(),
workspace_buffer=workspace_buffer,
)
self.mlp = Qwen3_VisionMLP(
dim,
@@ -197,6 +201,8 @@ class Qwen3_VisionBlock(nn.Module):
rotary_pos_emb_cos: torch.Tensor,
rotary_pos_emb_sin: torch.Tensor,
output_ws: Optional[torch.Tensor] = None,
max_seqlen: Optional[torch.Tensor] = None,
sequence_lengths: Optional[torch.Tensor] = None,
) -> torch.Tensor:
hidden_states = self.norm1(x)
hidden_states = rearrange(hidden_states, "s b ... -> b s ...")
@@ -206,6 +212,8 @@ class Qwen3_VisionBlock(nn.Module):
rotary_pos_emb_cos=rotary_pos_emb_cos,
rotary_pos_emb_sin=rotary_pos_emb_sin,
output_ws=output_ws,
max_seqlen=max_seqlen,
sequence_lengths=sequence_lengths,
)
attn = rearrange(attn, "b s ... -> s b ...")
x += attn
@@ -325,6 +333,18 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
is_neox_style=True,
)
workspace_buffer = None
if get_global_server_args().mm_attention_backend == "flashinfer_cudnn":
if torch.cuda.is_available() and (not is_npu()):
ws_device = torch.device("cuda", torch.cuda.current_device())
else:
ws_device = self.device
workspace_buffer = torch.empty(
FLASHINFER_WORKSPACE_SIZE_BYTES,
dtype=torch.uint8,
device=ws_device,
)
self.blocks = nn.ModuleList(
[
Qwen3_VisionBlock(
@@ -336,6 +356,7 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
quant_config=quant_config,
prefix=add_prefix(f"blocks.{layer_idx}", prefix),
use_data_parallel=use_data_parallel,
workspace_buffer=workspace_buffer,
)
for layer_idx in range(vision_config.depth)
]
@@ -468,6 +489,112 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
return torch.cat(result_parts, dim=0)
def _torch_interp_indices(
self, dim_size: int, device: torch.device
) -> torch.Tensor:
side = self.num_grid_per_side
if self.align_corners:
# align_corners=True
return torch.linspace(
0, side - 1, dim_size, dtype=torch.float32, device=device
)
else:
# align_corners=False (match _get_interpolation_indices)
idx = (torch.arange(dim_size, dtype=torch.float32, device=device) + 0.5) * (
side / dim_size
) - 0.5
return idx.clamp_(0, side - 1)
def fast_pos_embed_interpolate_from_list(self, grid_thw):
num_grid_per_side = self.num_grid_per_side
m_size = self.spatial_merge_size
hidden_dim = self.pos_embed.embedding_dim
outputs = []
for t, h, w in grid_thw:
h_idxs = torch.linspace(
0, num_grid_per_side - 1, h, dtype=torch.float32, device=self.device
)
w_idxs = torch.linspace(
0, num_grid_per_side - 1, w, dtype=torch.float32, device=self.device
)
h_floor = h_idxs.to(torch.long)
w_floor = w_idxs.to(torch.long)
h_ceil = torch.clamp(h_floor + 1, max=num_grid_per_side - 1)
w_ceil = torch.clamp(w_floor + 1, max=num_grid_per_side - 1)
dh = h_idxs - h_floor
dw = w_idxs - w_floor
# Create meshgrid view for all h, w vars
dh_grid, dw_grid = torch.meshgrid(dh, dw, indexing="ij")
h_floor_grid, w_floor_grid = torch.meshgrid(h_floor, w_floor, indexing="ij")
h_ceil_grid, w_ceil_grid = torch.meshgrid(h_ceil, w_ceil, indexing="ij")
# original computation of weights
# w00 = (1 - dh_grid) * (1 - dw_grid)
# w01 = (1 - dh_grid) * dw_grid
# w10 = dh_grid * (1 - dw_grid)
# w11 = dh_grid * dw_grid
# we reuse w11 here to avoid duplicate
# dh_grid * dw_grid computation
w11 = dh_grid * dw_grid
w10 = dh_grid - w11
w01 = dw_grid - w11
w00 = 1 - dh_grid - w01
h_grid = torch.stack([h_floor_grid, h_floor_grid, h_ceil_grid, h_ceil_grid])
w_grid = torch.stack([w_floor_grid, w_ceil_grid, w_floor_grid, w_ceil_grid])
h_grid_idx = h_grid * num_grid_per_side
indices = (h_grid_idx + w_grid).reshape(4, -1)
weights = torch.stack([w00, w01, w10, w11], dim=0).reshape(4, -1, 1)
weights = weights.to(dtype=self.dtype)
embeds = self.pos_embed(indices)
embeds *= weights
combined = embeds.sum(dim=0)
combined = combined.reshape(
h // m_size, m_size, w // m_size, m_size, hidden_dim
)
combined = combined.permute(0, 2, 1, 3, 4).reshape(1, -1, hidden_dim)
repeated = combined.expand(t, -1, -1).reshape(-1, hidden_dim)
outputs.append(repeated)
return torch.cat(outputs, dim=0)
def add_padding_to_fi_seqlens(
self, seq: np.ndarray, batch_size: int, padding_value: int
) -> np.ndarray:
batch_size_padded = next(
(b for b in BATCH_BUCKETS if b >= batch_size),
# For large batches (> max bucket), round up to a multiple of
# the base bucket size to avoid negative pad length.
round_up(batch_size, BATCH_BUCKETS[0]),
)
if batch_size_padded == batch_size:
return seq
return np.concatenate(
[
seq,
np.full(
(batch_size_padded - batch_size,), padding_value, dtype=seq.dtype
),
]
)
def bucket_flashinfer_max_seqlen(self, real_max_seqlen: int) -> int:
if real_max_seqlen <= 0:
return FLASHINFER_MAX_SEQLEN_BUCKETS[0]
return next(
(s for s in FLASHINFER_MAX_SEQLEN_BUCKETS if s >= real_max_seqlen),
# For large sequences (> max bucket), round up to a multiple of
# the largest bucket to avoid under-estimation.
round_up(real_max_seqlen, FLASHINFER_MAX_SEQLEN_BUCKETS[-1]),
)
def fast_pos_embed_interpolate(self, grid_thw):
"""Interpolate position embeddings for (batch, 3) size input dimensions.
@@ -523,6 +650,71 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
patch_pos_embeds, temporal_dims, height_dims, width_dims
)
def compute_flashinfer_batch_offsets_packed(
self,
token_cu_seqlens: np.ndarray,
*,
elem_per_token: int,
) -> np.ndarray:
"""
Build packed *element* indptrs for FlashInfer cuDNN prefill.
Input:
token_cu_seqlens: (B+1,) token indptr
elem_per_token: per-token element width on THIS TP rank
(usually hidden_size / attn_tp_size)
Output:
packed_offsets: (3 * (B_padded + 1),) int32
[qk_indptr, v_indptr, o_indptr] concatenated,
each indptr is (B_padded + 1,) in element units.
"""
assert token_cu_seqlens.ndim == 1 and token_cu_seqlens.size >= 2
B = int(token_cu_seqlens.size - 1)
B_padded = self.bucket_flashinfer_batch_size(B)
# token indptr -> pad to (B_padded+1,) by appending total_tokens for extra empty sequences
token_indptr = token_cu_seqlens.astype(np.int64, copy=False) # (B+1,)
if B_padded != B:
pad = np.full((B_padded - B,), token_indptr[-1], dtype=token_indptr.dtype)
token_indptr = np.concatenate([token_indptr, pad], axis=0) # (B_padded+1,)
# convert token indptr -> element indptr
elem_indptr = (token_indptr * int(elem_per_token)).astype(
np.int32
) # (B_padded+1,)
# q/k/v/o in this ViT path share the same indptr
return np.concatenate([elem_indptr, elem_indptr, elem_indptr], axis=0)
def bucket_flashinfer_batch_size(self, batch_size: int) -> int:
"""Bucketize batch size for cuDNN graph caching."""
return next(
(b for b in BATCH_BUCKETS if b >= batch_size),
round_up(batch_size, BATCH_BUCKETS[0]),
)
def compute_flashinfer_sequence_lengths_padded(
self,
token_cu_seqlens: np.ndarray,
) -> np.ndarray:
"""
token_cu_seqlens: (B+1,) token indptr
return: (B_padded,) token lengths (padded with 0)
"""
assert token_cu_seqlens.ndim == 1 and token_cu_seqlens.size >= 2
B = int(token_cu_seqlens.size - 1)
seq_lens = (token_cu_seqlens[1:] - token_cu_seqlens[:-1]).astype(
np.int32
) # (B,)
B_padded = self.bucket_flashinfer_batch_size(B)
if B_padded != B:
pad = np.zeros((B_padded - B,), dtype=np.int32)
seq_lens = np.concatenate([seq_lens, pad], axis=0) # (B_padded,)
return seq_lens
def forward(
self,
x: torch.Tensor,
@@ -536,24 +728,76 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
if isinstance(grid_thw, list):
grid_thw_list = grid_thw
grid_thw = torch.tensor(grid_thw, dtype=torch.int32)
grid_thw = np.array(grid_thw, dtype=np.int32)
else:
grid_thw_list = grid_thw.tolist()
grid_thw = grid_thw.cpu().numpy()
pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
pos_embeds = self.fast_pos_embed_interpolate_from_list(grid_thw_list)
x += pos_embeds
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)
# cu_seqlens must be on cpu because of npu_flash_attention_unpad operator restriction
if not is_npu():
cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)
# ---- build token indptr (B+1,) ----
token_cu_seqlens = np.repeat(
grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
).cumsum(axis=0, dtype=np.int32)
token_cu_seqlens = np.concatenate(
[np.zeros(1, dtype=np.int32), token_cu_seqlens]
)
flashinfer_max_seqlen = 0
cu_seqlens = None
if get_global_server_args().mm_attention_backend == "flashinfer_cudnn":
# real token lens (B,)
real_seq_lens = token_cu_seqlens[1:] - token_cu_seqlens[:-1]
flashinfer_max_seqlen = self.bucket_flashinfer_max_seqlen(
int(real_seq_lens.max()) if real_seq_lens.size > 0 else 0
)
# (B_padded,) token lengths
seq_lens_padded = self.compute_flashinfer_sequence_lengths_padded(
token_cu_seqlens
)
# element-per-token width on THIS ATTENTION TP rank
# q/k/v in VisionAttention are sharded by attention TP
attn_tp_size = 1 if self.use_data_parallel else self.tp_size
elem_per_token = (
self.hidden_size // attn_tp_size
) # == heads_per_rank * head_dim
# (3*(B_padded+1),) packed element indptrs
offsets_packed = self.compute_flashinfer_batch_offsets_packed(
token_cu_seqlens,
elem_per_token=elem_per_token,
)
sequence_lengths = (
torch.from_numpy(seq_lens_padded)
.to(device=self.device, dtype=torch.int32, non_blocking=True)
.view(-1, 1, 1, 1)
) # match cuDNN test style
cu_seqlens = torch.from_numpy(offsets_packed).to(
device=self.device, dtype=torch.int32, non_blocking=True
)
max_seqlen = int(flashinfer_max_seqlen)
sequence_lengths = sequence_lengths.to(self.device, non_blocking=True)
else:
cu_seqlens = cu_seqlens.to("cpu")
sequence_lengths = None
cu_seqlens = torch.from_numpy(token_cu_seqlens)
if not is_npu():
cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)
else:
cu_seqlens = cu_seqlens.to("cpu")
max_seqlen = None
x = x.unsqueeze(1)
cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)
deepstack_feature_lists = []
num_deepstack_captured = 0
@@ -563,6 +807,8 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
cu_seqlens=cu_seqlens,
rotary_pos_emb_cos=rotary_pos_emb_cos,
rotary_pos_emb_sin=rotary_pos_emb_sin,
max_seqlen=max_seqlen,
sequence_lengths=sequence_lengths,
)
if layer_num in self.deepstack_visual_indexes:

10
python/sglang/srt/server_args.py
View File

@@ -3867,7 +3867,15 @@ class ServerArgs:
parser.add_argument(
"--mm-attention-backend",
type=str,
choices=["sdpa", "fa3", "fa4", "triton_attn", "ascend_attn", "aiter_attn"],
choices=[
"sdpa",
"fa3",
"fa4",
"triton_attn",
"ascend_attn",
"aiter_attn",
"flashinfer_cudnn",
],
default=ServerArgs.mm_attention_backend,
help="Set multimodal attention backend.",
)