diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 303d41193..a8b25ecf6 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/python/sglang/srt/layers/attention/vision.py @@ -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, ) diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index d641826e3..685406fc6 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -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: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index a1bbd2038..e36a72eb7 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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.", ) diff --git a/test/manual/nightly/test_vlms_vit_flashinfer_cudnn.py b/test/manual/nightly/test_vlms_vit_flashinfer_cudnn.py new file mode 100644 index 000000000..1706bbda2 --- /dev/null +++ b/test/manual/nightly/test_vlms_vit_flashinfer_cudnn.py @@ -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]])