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

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
Co-authored-by: kousakawang <wanghanpei@bytedance.com>

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,10 +294,30 @@ class VisionTritonAttention(nn.Module):
         Returns:
              [b * s, h, head_size]
         """
+        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")
+
+            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(
@@ -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,8 +363,19 @@ class VisionFlash3Attention(nn.Module):
         Returns:
              [b * s, h, head_size]
         """
+        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()
@@ -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

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(

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,
+        )

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]])