[VLM] Introduce FlashInfer CUDNN Prefill as ViT Backend (#19003)

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

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

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
+        # ---- build token indptr (B+1,) ----
-        cu_seqlens = compute_cu_seqlens_from_grid_numpy(grid_thw)
+        token_cu_seqlens = np.repeat(
-        # cu_seqlens must be on cpu because of npu_flash_attention_unpad operator restriction
+            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:
+            sequence_lengths = None
+            cu_seqlens = torch.from_numpy(token_cu_seqlens)
             if not is_npu():
                 cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)
             else:
                 cu_seqlens = cu_seqlens.to("cpu")
+            max_seqlen = None
+
         x = x.unsqueeze(1)
 
+        cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)
+
         deepstack_feature_lists = []
         num_deepstack_captured = 0
 
@@ -563,6 +807,8 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
                 cu_seqlens=cu_seqlens,
                 rotary_pos_emb_cos=rotary_pos_emb_cos,
                 rotary_pos_emb_sin=rotary_pos_emb_sin,
+                max_seqlen=max_seqlen,
+                sequence_lengths=sequence_lengths,
             )
 
             if layer_num in self.deepstack_visual_indexes:

python/sglang/srt/server_args.py

@@ -3867,7 +3867,15 @@ class ServerArgs:
         parser.add_argument(
             "--mm-attention-backend",
             type=str,
-            choices=["sdpa", "fa3", "fa4", "triton_attn", "ascend_attn", "aiter_attn"],
+            choices=[
+                "sdpa",
+                "fa3",
+                "fa4",
+                "triton_attn",
+                "ascend_attn",
+                "aiter_attn",
+                "flashinfer_cudnn",
+            ],
             default=ServerArgs.mm_attention_backend,
             help="Set multimodal attention backend.",
         )

test/manual/nightly/test_vlms_vit_flashinfer_cudnn.py

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