chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 15 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
co-authored by luoyuan.luo
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
python/sglang/srt/layers/attention/vision.py
+152
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,
)
python/sglang/srt/models/qwen3_vl.py
+259 -13
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:
python/sglang/srt/server_args.py
+9 -1
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.",
)
test/manual/nightly/test_vlms_vit_flashinfer_cudnn.py
+258
View File
@@ -0,0 +1,258 @@
import argparse
import glob
import json
import os
import random
import sys
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.kits.mmmu_vlm_kit import _run_lmms_eval_with_retry
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_in_ci,
popen_launch_server,
)
MODELS = [
SimpleNamespace(model="Qwen/Qwen3-VL-30B-A3B-Instruct", mmmu_accuracy=0.51),
]
# Set default mem_fraction_static to 0.8
DEFAULT_MEM_FRACTION_STATIC = 0.8
class TestVLMViTFlashinferCudnn(CustomTestCase):
parsed_args = None # Class variable to store args
@classmethod
def setUpClass(cls):
# Removed argument parsing from here
cls.base_url = DEFAULT_URL_FOR_TEST
cls.api_key = "sk-123456"
cls.time_out = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
if cls.parsed_args is None:
cls.parsed_args = SimpleNamespace(
mem_fraction_static=DEFAULT_MEM_FRACTION_STATIC
)
# Set OpenAI API key and base URL environment variables. Needed for lmm-evals to work.
os.environ["OPENAI_API_KEY"] = cls.api_key
os.environ["OPENAI_API_BASE"] = f"{cls.base_url}/v1"
def run_mmmu_eval(
self,
model_version: str,
output_path: str,
*,
env: dict | None = None,
):
"""
Evaluate a VLM on the MMMU validation set with lmms‑eval.
Only `model_version` (checkpoint) and `chat_template` vary;
We are focusing only on the validation set due to resource constraints.
"""
# -------- fixed settings --------
model = "openai_compatible"
tp = 1
tasks = "mmmu_val"
batch_size = 32
log_suffix = "openai_compatible"
os.makedirs(output_path, exist_ok=True)
# -------- compose --model_args --------
model_args = f'model_version="{model_version}",' f"tp={tp}"
# -------- build command list --------
cmd = [
"python3",
"-m",
"lmms_eval",
"--model",
model,
"--model_args",
model_args,
"--tasks",
tasks,
"--batch_size",
str(batch_size),
"--output_path",
str(output_path),
]
_run_lmms_eval_with_retry(cmd, timeout=3600)
def _run_vlm_mmmu_test(
self,
model,
output_path,
test_name="",
custom_env=None,
log_level="info",
capture_output=False,
):
"""
Common method to run VLM MMMU benchmark test.
Args:
model: Model to test
output_path: Path for output logs
test_name: Optional test name for logging
custom_env: Optional custom environment variables
log_level: Log level for server (default: "info")
capture_output: Whether to capture server stdout/stderr
"""
print(f"\nTesting model: {model.model}{test_name}")
process = None
mmmu_accuracy = 0 # Initialize to handle potential exceptions
server_output = ""
try:
# Prepare environment variables
process_env = os.environ.copy()
if custom_env:
process_env.update(custom_env)
# if test vlm with cuda_ipc feature, open this env_var
process_env["SGLANG_USE_CUDA_IPC_TRANSPORT"] = "1"
# Prepare stdout/stderr redirection if needed
stdout_file = None
stderr_file = None
if capture_output:
stdout_file = open("/tmp/server_stdout.log", "w")
stderr_file = open("/tmp/server_stderr.log", "w")
# Launch server for testing
process = popen_launch_server(
model.model,
base_url=self.base_url,
timeout=self.time_out,
api_key=self.api_key,
other_args=[
"--mm-attention-backend",
"flashinfer_cudnn",
"--chunked-prefill-size",
"8192",
"--disable-radix-cache",
],
env=process_env,
return_stdout_stderr=(
(stdout_file, stderr_file) if capture_output else None
),
)
# Run evaluation
self.run_mmmu_eval(model.model, output_path)
# Get the result file
# Search recursively for JSON result files (lmms-eval v0.4.1+ creates subdirectories)
result_files = glob.glob(f"{output_path}/**/*.json", recursive=True)
if not result_files:
result_files = glob.glob(f"{output_path}/*.json")
if not result_files:
raise FileNotFoundError(f"No JSON result files found in {output_path}")
result_file_path = result_files[0]
with open(result_file_path, "r") as f:
result = json.load(f)
print(f"Result{test_name}\n: {result}")
# Process the result
mmmu_accuracy = result["results"]["mmmu_val"]["mmmu_acc,none"]
print(
f"Model {model.model} achieved accuracy{test_name}: {mmmu_accuracy:.4f}"
)
# Capture server output if requested
if capture_output and process:
server_output = self._read_output_from_files()
# Assert performance meets expected threshold
self.assertGreaterEqual(
mmmu_accuracy,
model.mmmu_accuracy,
f"Model {model.model} accuracy ({mmmu_accuracy:.4f}) below expected threshold ({model.mmmu_accuracy:.4f}){test_name}",
)
return server_output
except Exception as e:
print(f"Error testing {model.model}{test_name}: {e}")
self.fail(f"Test failed for {model.model}{test_name}: {e}")
finally:
# Ensure process cleanup happens regardless of success/failure
if process is not None and process.poll() is None:
print(f"Cleaning up process {process.pid}")
try:
kill_process_tree(process.pid)
except Exception as e:
print(f"Error killing process: {e}")
# clean up temporary files
if capture_output:
if stdout_file:
stdout_file.close()
if stderr_file:
stderr_file.close()
for filename in ["/tmp/server_stdout.log", "/tmp/server_stderr.log"]:
try:
if os.path.exists(filename):
os.remove(filename)
except Exception as e:
print(f"Error removing {filename}: {e}")
def _read_output_from_files(self):
output_lines = []
log_files = [
("/tmp/server_stdout.log", "[STDOUT]"),
("/tmp/server_stderr.log", "[STDERR]"),
]
for filename, tag in log_files:
try:
if os.path.exists(filename):
with open(filename, "r") as f:
for line in f:
output_lines.append(f"{tag} {line.rstrip()}")
except Exception as e:
print(f"Error reading {tag.lower()} file: {e}")
return "\n".join(output_lines)
def test_vlm_mmmu_benchmark(self):
"""Test VLM models against MMMU benchmark."""
models_to_test = MODELS
if is_in_ci():
models_to_test = [random.choice(MODELS)]
for model in models_to_test:
self._run_vlm_mmmu_test(model, "./logs")
if __name__ == "__main__":
# Define and parse arguments here, before unittest.main
parser = argparse.ArgumentParser(description="Test VLM models")
parser.add_argument(
"--mem-fraction-static",
type=float,
help="Static memory fraction for the model",
default=DEFAULT_MEM_FRACTION_STATIC,
)
# Parse args intended for unittest
args = parser.parse_args()
# Store the parsed args object on the class
TestVLMViTFlashinferCudnn.parsed_args = args
# Pass args to unittest
unittest.main(argv=[sys.argv[0]])