chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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,
)

271
test/manual/nightly/test_vlms_vit_cuda_graph.py Normal file
View File

@@ -0,0 +1,271 @@
import argparse
import glob
import json
import os
import random
import subprocess
import sys
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
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/Qwen2.5-VL-7B-Instruct", mmmu_accuracy=0.60),
]
# Set default mem_fraction_static to 0.8
DEFAULT_MEM_FRACTION_STATIC = 0.8
class TestVLMViTCudaGraph(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
cls.enable_vit_cuda_graph = "1"
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"
os.environ["SGLANG_VIT_ENABLE_CUDA_GRAPH"] = cls.enable_vit_cuda_graph
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),
]
subprocess.run(
cmd,
check=True,
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"
process_env["SGLANG_VIT_ENABLE_CUDA_GRAPH"] = "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",
"fa3",
"--enable-piecewise-cuda-graph",
"--piecewise-cuda-graph-max-tokens",
"8192",
"--chunked-prefill-size",
"8192",
"--disable-radix-cache",
"--disable-overlap-schedule",
"--piecewise-cuda-graph-compiler",
"eager",
],
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
TestVLMViTCudaGraph.parsed_args = args
# Pass args to unittest
unittest.main(argv=[sys.argv[0]])