From 3b8a824b8b2ed140e0ed97f7227eda2349818576 Mon Sep 17 00:00:00 2001 From: Yuan Luo Date: Sat, 13 Dec 2025 20:49:40 +0800 Subject: [PATCH] [VLM] Support VLM ViT Piecewise CUDA Graph (#14422) Co-authored-by: luoyuan.luo Co-authored-by: kousakawang --- python/sglang/srt/layers/attention/vision.py | 103 +++++-- python/sglang/srt/models/qwen2_5_vl.py | 84 +++++- .../srt/multimodal/vit_cuda_graph_runner.py | 263 +++++++++++++++++ .../nightly/test_vlms_vit_cuda_graph.py | 271 ++++++++++++++++++ 4 files changed, 690 insertions(+), 31 deletions(-) create mode 100644 python/sglang/srt/multimodal/vit_cuda_graph_runner.py create mode 100644 test/manual/nightly/test_vlms_vit_cuda_graph.py diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 9f556a885..a3f2d3080 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/python/sglang/srt/layers/attention/vision.py @@ -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 diff --git a/python/sglang/srt/models/qwen2_5_vl.py b/python/sglang/srt/models/qwen2_5_vl.py index b84b17b3e..7c1d09852 100644 --- a/python/sglang/srt/models/qwen2_5_vl.py +++ b/python/sglang/srt/models/qwen2_5_vl.py @@ -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( diff --git a/python/sglang/srt/multimodal/vit_cuda_graph_runner.py b/python/sglang/srt/multimodal/vit_cuda_graph_runner.py new file mode 100644 index 000000000..f5ba1b86d --- /dev/null +++ b/python/sglang/srt/multimodal/vit_cuda_graph_runner.py @@ -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, + ) diff --git a/test/manual/nightly/test_vlms_vit_cuda_graph.py b/test/manual/nightly/test_vlms_vit_cuda_graph.py new file mode 100644 index 000000000..32948e398 --- /dev/null +++ b/test/manual/nightly/test_vlms_vit_cuda_graph.py @@ -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]])