vlm: batch cross-request vit encoding and reuse attention metadata (#24013)

Co-authored-by: yhyang201 <yhyang201@gmail.com>
Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/srt/layers/attention/intel_amx_backend.py

@@ -212,6 +212,9 @@ class IntelAMXAttnBackend(AttentionBackend):
         seq_lens, extend_seq_lens, extend_start_loc, tree_mask = self.extend_metadata
 
         _, max_extend_len = self.forward_metadata
+        seq_lens = forward_batch.seq_lens
+        if seq_lens.dtype != torch.int64:
+            seq_lens = seq_lens.to(torch.int64)
         self.extend_attention_fwd(
             q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
             k,
@@ -255,6 +258,9 @@ class IntelAMXAttnBackend(AttentionBackend):
             seq_lens = forward_batch.seq_lens
 
         q = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim)
+        seq_lens = forward_batch.seq_lens
+        if seq_lens.dtype != torch.int64:
+            seq_lens = seq_lens.to(torch.int64)
 
         if layer.qk_head_dim != layer.v_head_dim:
             o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))

python/sglang/srt/layers/attention/vision.py

@@ -46,10 +46,18 @@ _is_xpu = is_xpu()
 if _is_cuda:
     from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
 
-    from sglang.jit_kernel.flash_attention import (
+    from sglang.jit_kernel.flash_attention import flash_attn_varlen_func
-        flash_attn_varlen_func,
+
+    def flash_attn_func(*args, ver: int = 3, **kwargs):
+        if ver == 4:
+            from sglang.jit_kernel.flash_attention_v4 import (
+                flash_attn_varlen_func as flash_attn_varlen_func_fa4,
             )
 
+            return flash_attn_varlen_func_fa4(*args, **kwargs)
+        return flash_attn_varlen_func(*args, **kwargs)
+
+
 if _is_cpu and _is_cpu_amx_available:
     flash_attn_varlen_func = torch.ops.sgl_kernel.flash_attn_varlen_func
 
@@ -78,7 +86,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.runtime_context import get_server_args
-from sglang.srt.utils import add_prefix, get_bool_env_var
+from sglang.srt.utils import add_prefix
 
 _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
 
@@ -121,6 +129,42 @@ class SingletonCache:
         return self.get_data() is None
 
 
+@dataclasses.dataclass
+class VisionAttentionMetadata:
+
+    cu_seqlens: torch.Tensor
+    seq_lens: torch.Tensor
+    max_seqlen: int
+
+    # flashinfer_cudnn specific (optional)
+    packed_indptrs: Optional[torch.Tensor] = None
+    sequence_lengths: Optional[torch.Tensor] = None
+    flashinfer_max_seqlen: Optional[int] = None
+
+
+def prepare_vision_attention_metadata(
+    cu_seqlens: torch.Tensor,
+    device: torch.device,
+    *,
+    packed_indptrs: Optional[torch.Tensor] = None,
+    sequence_lengths: Optional[torch.Tensor] = None,
+    flashinfer_max_seqlen: Optional[int] = None,
+) -> VisionAttentionMetadata:
+    # Compute all attention metadata once before the encoder layer loop.
+
+    cu_seqlens = cu_seqlens.to(device=device, dtype=torch.int32, non_blocking=True)
+    seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+    max_seqlen = int(seq_lens.max().item())
+    return VisionAttentionMetadata(
+        cu_seqlens=cu_seqlens,
+        seq_lens=seq_lens,
+        max_seqlen=max_seqlen,
+        packed_indptrs=packed_indptrs,
+        sequence_lengths=sequence_lengths,
+        flashinfer_max_seqlen=flashinfer_max_seqlen,
+    )
+
+
 # TODO: requires real seqlens from images
 @functools.lru_cache(maxsize=128)
 def _get_cu_seqlens_for_shape(batch_size: int, seqlen: int, device) -> torch.Tensor:
@@ -205,7 +249,7 @@ class VisionSdpaAttention(nn.Module):
         dropout: float = 0.0,
         flatten_batch: bool = False,
         softmax_in_single_precision: bool = False,
-        softmax_scale: float | None = None,
+        softmax_scale: Optional[float] = None,
         **kwargs,
     ):
         super().__init__()
@@ -285,7 +329,6 @@ class VisionSdpaAttention(nn.Module):
         bsz: int,
         cu_seqlens: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
-        softmax_scale: Optional[float] = None,
         **kwargs,
     ) -> torch.Tensor:
         r"""
@@ -371,20 +414,25 @@ class VisionTritonAttention(nn.Module):
         q: torch.Tensor,
         k: torch.Tensor,
         v: torch.Tensor,
-        cu_seqlens: torch.Tensor | SingletonCache | None,
+        cu_seqlens: torch.Tensor | SingletonCache | list | None,
         bsz: int,
         seq_len: int,
         softmax_scale: Optional[float] = None,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
         **kwargs,
     ) -> torch.Tensor:
         r"""
         Args:
             cu_seqlens: [b]
-            softmax_scale: override softmax scale (default 1/sqrt(head_dim))
         Returns:
              [b * s, h, head_size]
         """
-        if envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get():
+        if forward_metadata is not None:
+            cu_seqlens_gpu = forward_metadata.cu_seqlens
+            seq_lens = forward_metadata.seq_lens
+            max_seqlen = forward_metadata.max_seqlen
+            output = torch.empty_like(q)
+        elif envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get():
             if "output_ws" not in kwargs:
                 raise RuntimeError("output_ws should be prepared for cuda-graph mode")
 
@@ -392,37 +440,28 @@ class VisionTritonAttention(nn.Module):
                 raise RuntimeError("cuda-graph mode cu_seqlens should be a list")
 
             output = kwargs["output_ws"]
-            context_attention_fwd(
+            cu_seqlens_gpu = cu_seqlens[0]
-                q,
+            seq_lens = cu_seqlens[1]
-                k,
+            max_seqlen = cu_seqlens[2]
-                v,
-                output,
-                cu_seqlens[0],
-                cu_seqlens[1],
-                cu_seqlens[2],
-                is_causal=False,
-                sm_scale=softmax_scale,
-            )
         else:
-            cu_seqlens = resolve_seqlens(cu_seqlens, bsz, seq_len, device=q.device)
+            cu_seqlens_gpu = resolve_seqlens(cu_seqlens, bsz, seq_len, device=q.device)
-
-            # [b * s, head, head_size]
-            output = torch.empty_like(q)
-
             seq_lens = kwargs.get("sequence_lengths")
             if seq_lens is None:
-                seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+                seq_lens = cu_seqlens_gpu[1:] - cu_seqlens_gpu[:-1]
             else:
                 seq_lens = seq_lens.to(device=q.device, dtype=torch.int32)
             max_seqlen = resolve_precomputed_max_seqlen(
-                cu_seqlens, kwargs.get("max_seqlen")
+                cu_seqlens_gpu, kwargs.get("max_seqlen")
             )
+            # [b * s, head, head_size]
+            output = torch.empty_like(q)
+
         context_attention_fwd(
             q,
             k,
             v,
             output,
-                cu_seqlens.to(q.device),
+            cu_seqlens_gpu,
             seq_lens,
             max_seqlen,
             is_causal=False,
@@ -450,10 +489,11 @@ class VisionFlash3Attention(nn.Module):
         q: torch.Tensor,
         k: torch.Tensor,
         v: torch.Tensor,
-        cu_seqlens: torch.Tensor | SingletonCache | None,
+        cu_seqlens: torch.Tensor | SingletonCache | list | None,
         bsz: int,
         seq_len: int,
         softmax_scale: Optional[float] = None,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
         **kwargs,
     ) -> torch.Tensor:
         r"""
@@ -465,29 +505,24 @@ class VisionFlash3Attention(nn.Module):
         window_size = kwargs.get("window_size", (-1, -1))
         s_aux = kwargs.get("s_aux", None)
 
-        if envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get():
+        if forward_metadata is not None:
+            cu_seqlens_gpu = forward_metadata.cu_seqlens
+            max_seqlen = forward_metadata.max_seqlen
+        elif envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get():
+            if not isinstance(cu_seqlens, list):
+                raise RuntimeError("cuda-graph mode cu_seqlens should be a list")
+            cu_seqlens_gpu = cu_seqlens[0]
             max_seqlen = cu_seqlens[1]
-            fa_kwargs = dict(
-                cu_seqlens_q=cu_seqlens[0],
-                cu_seqlens_k=cu_seqlens[0],
-                max_seqlen_q=max_seqlen,
-                max_seqlen_k=max_seqlen,
-                softmax_scale=softmax_scale,
-                window_size=window_size,
-            )
-            if s_aux is not None:
-                fa_kwargs["sinks"] = s_aux
-            output = flash_attn_varlen_func(q, k, v, **fa_kwargs)
         else:
-            cu_seqlens = resolve_seqlens(cu_seqlens, bsz, seq_len, device=q.device)
+            cu_seqlens_gpu = resolve_seqlens(cu_seqlens, bsz, seq_len, device=q.device)
-            cu_seqlens = cu_seqlens.to(dtype=torch.int32).to(q.device)
+            cu_seqlens_gpu = cu_seqlens_gpu.to(dtype=torch.int32).to(q.device)
             max_seqlen = resolve_precomputed_max_seqlen(
-                cu_seqlens, kwargs.get("max_seqlen")
+                cu_seqlens_gpu, kwargs.get("max_seqlen")
             )
 
         fa_kwargs = dict(
-                cu_seqlens_q=cu_seqlens,
+            cu_seqlens_q=cu_seqlens_gpu,
-                cu_seqlens_k=cu_seqlens,
+            cu_seqlens_k=cu_seqlens_gpu,
             max_seqlen_q=max_seqlen,
             max_seqlen_k=max_seqlen,
             softmax_scale=softmax_scale,
@@ -495,7 +530,7 @@ class VisionFlash3Attention(nn.Module):
         )
         if s_aux is not None:
             fa_kwargs["sinks"] = s_aux
-            output = flash_attn_varlen_func(q, k, v, **fa_kwargs)
+        output = flash_attn_func(q, k, v, **fa_kwargs)
 
         return output
 
@@ -518,6 +553,7 @@ class VisionFlash4Attention(nn.Module):
         bsz: int,
         seq_len: int,
         softmax_scale: Optional[float] = None,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
         **kwargs,
     ) -> torch.Tensor:
         r"""
@@ -526,26 +562,22 @@ class VisionFlash4Attention(nn.Module):
         Returns:
              [b * s, h, head_size]
         """
-        if cu_seqlens is None:
+        if forward_metadata is not None:
-            cu_seqlens = _get_cu_seqlens_for_shape(bsz, seq_len, device=q.device)
+            cu_seqlens_gpu = forward_metadata.cu_seqlens
-        elif isinstance(cu_seqlens, SingletonCache):
+            max_seqlen = forward_metadata.max_seqlen
-            if cu_seqlens.empty():
+        else:
-                cu_seqlens.set_data(
+            cu_seqlens_gpu = resolve_seqlens(cu_seqlens, bsz, seq_len, device=q.device)
-                    _get_cu_seqlens_for_shape(bsz, seq_len, device=q.device)
+            cu_seqlens_gpu = cu_seqlens_gpu.to(dtype=torch.int32).to(q.device)
-                )
-            cu_seqlens = cu_seqlens.get_data()
-
-        cu_seqlens = cu_seqlens.to(dtype=torch.int32).to(q.device)
             max_seqlen = resolve_precomputed_max_seqlen(
-            cu_seqlens, kwargs.get("max_seqlen")
+                cu_seqlens_gpu, kwargs.get("max_seqlen")
             )
 
-        output = flash_attn_varlen_func(
+        output = flash_attn_func(
             q,
             k,
             v,
-            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_q=cu_seqlens_gpu,
-            cu_seqlens_k=cu_seqlens,
+            cu_seqlens_k=cu_seqlens_gpu,
             max_seqlen_q=max_seqlen,
             max_seqlen_k=max_seqlen,
             softmax_scale=softmax_scale,
@@ -576,6 +608,7 @@ class VisionFlashInferAttention(nn.Module):
         bsz: int,
         seq_len: int,
         softmax_scale: Optional[float] = None,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
         **kwargs,
     ) -> torch.Tensor:
         r"""
@@ -584,6 +617,12 @@ class VisionFlashInferAttention(nn.Module):
         Returns:
              [b * s, h, head_size]
         """
+        # ---- resolve sequence_lengths, packed indptrs, max_seqlen ----
+        if forward_metadata is not None and forward_metadata.packed_indptrs is not None:
+            sequence_lengths = forward_metadata.sequence_lengths
+            packed_cu_seqlens = forward_metadata.packed_indptrs
+            max_seqlen = forward_metadata.flashinfer_max_seqlen
+        else:
             if "sequence_lengths" not in kwargs:
                 raise RuntimeError(
                     "sequence_lengths should be prepared for vision flashinfer_cudnn attention backend"
@@ -592,8 +631,8 @@ class VisionFlashInferAttention(nn.Module):
                 raise RuntimeError(
                     "max_seqlen should be prepared for vision flashinfer_cudnn attention backend"
                 )
-
+            sequence_lengths = kwargs["sequence_lengths"]
-        sequence_lengths = kwargs["sequence_lengths"]  # (B_padded,) or (B_padded,1,1,1)
+            packed_cu_seqlens = cu_seqlens
             max_seqlen = kwargs["max_seqlen"]
 
         # max_seqlen must be python int
@@ -611,7 +650,7 @@ class VisionFlashInferAttention(nn.Module):
             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):
+        if not isinstance(packed_cu_seqlens, torch.Tensor):
             raise RuntimeError(
                 "flashinfer_cudnn expects packed indptrs as a torch.Tensor"
             )
@@ -624,7 +663,9 @@ class VisionFlashInferAttention(nn.Module):
 
         # 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)
+        cu_seqlens_1d = packed_cu_seqlens.view(-1).to(
+            device=q.device, dtype=torch.int32
+        )
         expected = 3 * (B + 1)
         if int(cu_seqlens_1d.numel()) != expected:
             raise RuntimeError(
@@ -704,20 +745,24 @@ class VisionAiterAttention(nn.Module):
         bsz: int,
         seq_len: int,
         softmax_scale: Optional[float] = None,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
         **kwargs,
     ) -> torch.Tensor:
-        cu_seqlens = resolve_seqlens(cu_seqlens, bsz, seq_len, device=q.device)
+        if forward_metadata is not None:
-
+            cu_seqlens_gpu = forward_metadata.cu_seqlens
-        cu_seqlens = cu_seqlens.to(dtype=torch.int32).to(q.device)
+            max_seqlen = forward_metadata.max_seqlen
-        seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+        else:
+            cu_seqlens_gpu = resolve_seqlens(cu_seqlens, bsz, seq_len, device=q.device)
+            cu_seqlens_gpu = cu_seqlens_gpu.to(dtype=torch.int32).to(q.device)
+            seq_lens = cu_seqlens_gpu[1:] - cu_seqlens_gpu[:-1]
             max_seqlen = seq_lens.max().item()
 
         return self.flash_attn_varlen_func(
             q=q,
             k=k,
             v=v,
-            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_q=cu_seqlens_gpu,
-            cu_seqlens_k=cu_seqlens,
+            cu_seqlens_k=cu_seqlens_gpu,
             max_seqlen_q=max_seqlen,
             max_seqlen_k=max_seqlen,
             softmax_scale=softmax_scale,
@@ -743,6 +788,7 @@ class VisionAscendAttention(nn.Module):
         bsz: int,
         seq_len: int,
         softmax_scale: Optional[float] = None,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
         **kwargs,
     ) -> torch.Tensor:
         r"""
@@ -751,7 +797,13 @@ class VisionAscendAttention(nn.Module):
         Returns:
              [b * s, h, head_size]
         """
-        if envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get():
+        if forward_metadata is not None:
+            seq_lens = forward_metadata.seq_lens
+            if seq_lens.is_npu:
+                seq_lens = seq_lens.to("cpu")
+            output = torch.empty_like(q)
+            seq_len_arg = seq_lens.to(torch.int32)
+        elif envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get():
             if "output_ws" not in kwargs:
                 raise RuntimeError("output_ws should be prepared for npu-graph mode")
             output = kwargs["output_ws"]
@@ -955,6 +1007,7 @@ class VisionAttention(nn.Module):
         self.dropout = dropout
         num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
         self.head_size = head_dim if head_dim is not None else embed_dim // num_heads
+        self.softmax_scale = softmax_scale
         self.hidden_size_per_attention_head = dist_utils.divide(
             projection_size, num_heads
         )
@@ -994,7 +1047,6 @@ class VisionAttention(nn.Module):
         self.customized_position_embedding_applier = (
             customized_position_embedding_applier
         )
-        self.softmax_scale = softmax_scale
         self.qkv_backend = QKV_BACKEND_IMPL[qkv_backend](
             head_dim=self.head_size,
             num_heads=self.num_attention_heads_per_partition,
@@ -1185,6 +1237,7 @@ class VisionAttention(nn.Module):
         rotary_pos_emb_cos: Optional[torch.Tensor] = None,
         rotary_pos_emb_sin: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
         full_attn: bool = True,
         **kwargs,
     ) -> torch.Tensor:
@@ -1333,6 +1386,7 @@ class VisionAttention(nn.Module):
             seq_len=s,
             cu_seqlens=cu_seqlens,
             attention_mask=attention_mask,
+            forward_metadata=forward_metadata,
             sequence_lengths=sequence_lengths,
             max_seqlen=max_seqlen,
             output_ws=attn_output_ws,

python/sglang/srt/managers/mm_utils.py

@@ -9,6 +9,7 @@ import pickle
 import sys
 from abc import abstractmethod
 from collections import defaultdict
+from dataclasses import dataclass, field
 from multiprocessing import shared_memory
 from typing import Any, Callable, Dict, List, Literal, Optional, Tuple
 
@@ -33,10 +34,11 @@ from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalSta
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.multimodal.evs import EVSEmbeddingResult
 from sglang.srt.runtime_context import get_parallel, get_server_args
-from sglang.srt.utils import flatten_nested_list, is_npu, print_warning_once
+from sglang.srt.utils import flatten_nested_list, is_hip, is_npu, print_warning_once
 from sglang.srt.utils.stale_shm_cleanup import make_shm_name
 from sglang.utils import logger
 
+_is_hip = is_hip()
 _is_npu = is_npu()
 
 # NOTE: Using the shared logger from sglang.utils instead of creating a module-specific logger
@@ -581,6 +583,80 @@ def _get_chunked_embedding_full(
     return embedding_per_req_chunk, input_ids
 
 
+@dataclass
+class PerImageRequestInfo:
+    """Metadata for a single request using the per-image encoding path."""
+
+    req_idx: int
+    items: List[MultimodalDataItem]
+    items_offset: List[Tuple[int, int]]
+    extend_prefix_len: int
+    extend_seq_len: int
+    overlapping: List[Tuple[int, MultimodalDataItem, int, int]] = field(
+        default_factory=list
+    )
+
+
+def _batch_encode_per_image_misses(
+    data_embedding_func: DataEmbeddingFunc,
+    per_image_requests: List[PerImageRequestInfo],
+    device: torch.device,
+) -> Dict[int, torch.Tensor]:
+    """
+    Collect cache misses across ALL per-image requests, deduplicate by hash,
+    encode in a single ViT call, and populate the cache.
+
+    Returns:
+        hash_to_embedding: mapping from item.hash to its full embedding tensor.
+    """
+    unique_misses: Dict[int, Tuple[MultimodalDataItem, int]] = {}
+    hash_to_embedding: Dict[int, torch.Tensor] = {}
+
+    # Phase 1a: find overlapping items per request and collect cache misses
+    for req_info in per_image_requests:
+        chunk_start = req_info.extend_prefix_len
+        chunk_end = chunk_start + req_info.extend_seq_len  # exclusive
+        overlapping = []
+        if req_info.extend_seq_len > 0:
+            for idx, (item, (start, end)) in enumerate(
+                zip(req_info.items, req_info.items_offset)
+            ):
+                if end >= chunk_start and start < chunk_end:
+                    overlapping.append((idx, item, start, end))
+        req_info.overlapping = overlapping
+
+        for _idx, item, start, end in overlapping:
+            if item.hash in hash_to_embedding:
+                continue
+            cached = embedding_cache.get_single(item.hash)
+            if cached is not None:
+                hash_to_embedding[item.hash] = cached.embedding
+            elif item.hash not in unique_misses:
+                token_count = end - start + 1
+                unique_misses[item.hash] = (item, token_count)
+
+    # Phase 1b: single ViT call for all unique cache misses
+    if unique_misses:
+        ordered_hashes = list(unique_misses.keys())
+        miss_items = [unique_misses[h][0] for h in ordered_hashes]
+        token_counts = [unique_misses[h][1] for h in ordered_hashes]
+
+        if not _can_skip_pre_embed_feature_move(data_embedding_func):
+            _move_items_to_device(miss_items, device)
+        all_miss_embedding = data_embedding_func(miss_items)
+        all_miss_embedding = all_miss_embedding.reshape(
+            -1, all_miss_embedding.shape[-1]
+        )
+
+        split_embeddings = torch.split(all_miss_embedding, token_counts, dim=0)
+        for h, emb in zip(ordered_hashes, split_embeddings):
+            embedding_cache.set(h, EmbeddingResult(embedding=emb))
+            # Keep a local ref (no extra GPU memory) so assembly never fails due to LRU eviction.
+            hash_to_embedding[h] = emb
+
+    return hash_to_embedding
+
+
 def _get_chunked_embedding_by_item(
     data_embedding_func: DataEmbeddingFunc,
     embedding_items_per_req: List[MultimodalDataItem],
@@ -590,8 +666,7 @@ def _get_chunked_embedding_by_item(
     device: torch.device,
 ) -> Optional[torch.Tensor]:
     """
-    Per-image chunk-aware encoding: only encode images overlapping with the
+    Per-image chunk-aware encoding for one request.
-    current chunk, cache each image individually.
     Items must already be split per-image (each item has exactly one offset).
     """
     chunk_start = extend_prefix_len
@@ -600,20 +675,18 @@ def _get_chunked_embedding_by_item(
     if extend_seq_len <= 0:
         return None
 
-    # 1. Find items overlapping with current chunk
-    # offsets are (start, end) inclusive on both ends
     overlapping = []
-    for idx, (item, offset) in enumerate(zip(embedding_items_per_req, items_offset)):
+    for idx, (item, (start, end)) in enumerate(
-        start, end = offset
+        zip(embedding_items_per_req, items_offset)
+    ):
         if end >= chunk_start and start < chunk_end:
             overlapping.append((idx, item, start, end))
 
     if not overlapping:
         return None
 
-    # 2. Check per-image cache for each overlapping item
+    cached_embeddings = {}
-    cached_embeddings = {}  # idx -> tensor
+    miss_items = []
-    miss_items = []  # (idx, item, start, end)
     for idx, item, start, end in overlapping:
         cached = embedding_cache.get_single(item.hash)
         if cached is not None:
@@ -622,7 +695,6 @@ def _get_chunked_embedding_by_item(
         else:
             miss_items.append((idx, item, start, end))
 
-    # 3. Batch encode all cache-miss items in one ViT call
     if miss_items:
         miss_item_list = [item for _, item, _, _ in miss_items]
         if not _can_skip_pre_embed_feature_move(data_embedding_func):
@@ -632,19 +704,44 @@ def _get_chunked_embedding_by_item(
             -1, all_miss_embedding.shape[-1]
         )
 
-        # Split output by per-item token count
         token_counts = [end - start + 1 for _, _, start, end in miss_items]
         split_embeddings = torch.split(all_miss_embedding, token_counts, dim=0)
 
         for (idx, item, _, _), emb in zip(miss_items, split_embeddings):
             cached_embeddings[idx] = emb
-            emb_result = EmbeddingResult(embedding=emb)
+            embedding_cache.set(item.hash, EmbeddingResult(embedding=emb))
-            embedding_cache.set(item.hash, emb_result)
 
-    # 4. Assemble chunk: for each overlapping item, extract the overlap slice
     chunk_slices = []
     for idx, _, start, end in overlapping:
-        emb = cached_embeddings[idx]  # shape: (end - start + 1, hidden)
+        emb = cached_embeddings[idx]
+        overlap_start = max(start, chunk_start)
+        overlap_end = min(end, chunk_end - 1)  # inclusive
+        local_start = overlap_start - start
+        local_end = overlap_end - start + 1  # exclusive for slicing
+        chunk_slices.append(emb[local_start:local_end])
+
+    return torch.cat(chunk_slices, dim=0)
+
+
+def _assemble_per_image_chunk(
+    overlapping: List[Tuple[int, MultimodalDataItem, int, int]],
+    hash_to_embedding: Dict[int, torch.Tensor],
+    extend_prefix_len: int,
+    extend_seq_len: int,
+) -> Optional[torch.Tensor]:
+    """
+    Assemble the chunk embedding for one request from pre-computed embeddings.
+    All overlapping items must already have their embeddings in hash_to_embedding.
+    """
+    if not overlapping:
+        return None
+
+    chunk_start = extend_prefix_len
+    chunk_end = extend_prefix_len + extend_seq_len  # exclusive
+
+    chunk_slices = []
+    for _idx, item, start, end in overlapping:
+        emb = hash_to_embedding[item.hash]  # shape: (end - start + 1, hidden)
         overlap_start = max(start, chunk_start)
         overlap_end = min(end, chunk_end - 1)  # inclusive
         local_start = overlap_start - start
@@ -664,14 +761,19 @@ def _get_chunked_prefill_embedding(
     input_ids: torch.Tensor,
 ) -> tuple[torch.Tensor | None, torch.Tensor]:
     """
-    Chunked prefill embedding: encode per-request items and extract the chunk.
+    Chunked prefill embedding: encode items across all requests and extract
-    Items are already split per-image at processor stage.
+    per-request chunks. Images from all requests are batched into a single
+    ViT call for efficiency.
     """
-    embedding_list = []
     device = input_ids.device
     # FIXME(Xinyuan): temporary workaround for eagle3
     max_iterations = min(len(items_size) - 1, len(prefix_length))
 
+    # Phase 0: classify requests into per-image vs full/EVS path
+    per_image_requests = []  # batched ViT encoding
+    full_path_requests = []  # per-request encoding (EVS etc.)
+    all_chunks: List[Tuple[int, torch.Tensor]] = []
+
     for i in range(max_iterations):
         if items_size[i] == items_size[i + 1]:
             continue
@@ -681,18 +783,27 @@ def _get_chunked_prefill_embedding(
 
         extend_prefix_len = prefix_length[i]
         extend_seq_len = extend_length[i] if i < len(extend_length) else 0
+        if extend_seq_len <= 0:
+            continue
 
-        # Skip if all items already prefilled
+        # Skip if all items already prefilled.
         if all(offset_end < prefix_length[i] for _, offset_end in items_offset):
             continue
 
-        # Use per-image path when all items have exactly one offset (already
+        req_info = PerImageRequestInfo(
-        # split per-image) — this avoids encoding images not in this chunk.
+            req_idx=i,
-        # Fall back to combined path for non-split items or EVS.
+            items=embedding_items_per_req,
-        is_per_image = all(len(item.offsets) == 1 for item in embedding_items_per_req)
+            items_offset=items_offset,
+            extend_prefix_len=extend_prefix_len,
+            extend_seq_len=extend_seq_len,
+        )
 
+        is_per_image = all(len(item.offsets) == 1 for item in embedding_items_per_req)
         if is_per_image:
-            chunk_embedding = _get_chunked_embedding_by_item(
+            if _is_hip:
+                # ROCm CI regressed with one large cross-request ViT batch; keep
+                # the previous per-request path on HIP while CUDA uses batching.
+                chunk = _get_chunked_embedding_by_item(
                     data_embedding_func,
                     embedding_items_per_req,
                     items_offset,
@@ -700,20 +811,47 @@ def _get_chunked_prefill_embedding(
                     extend_seq_len,
                     device,
                 )
-            if chunk_embedding is not None:
+                if chunk is not None:
-                embedding_list.append(chunk_embedding)
+                    all_chunks.append((i, chunk))
             else:
+                per_image_requests.append(req_info)
+        else:
+            full_path_requests.append(req_info)
+
+    # Phase 1: batch encode all per-image cache misses in ONE ViT call
+    hash_to_embedding: Dict[int, torch.Tensor] = {}
+    if per_image_requests:
+        hash_to_embedding = _batch_encode_per_image_misses(
+            data_embedding_func, per_image_requests, device
+        )
+
+    # Phase 2: assemble per-request chunks in original request order
+    for req_info in per_image_requests:
+        chunk = _assemble_per_image_chunk(
+            req_info.overlapping,
+            hash_to_embedding,
+            req_info.extend_prefix_len,
+            req_info.extend_seq_len,
+        )
+        if chunk is not None:
+            all_chunks.append((req_info.req_idx, chunk))
+
+    for req_info in full_path_requests:
         chunk_embedding, input_ids = _get_chunked_embedding_full(
             data_embedding_func,
-                embedding_items_per_req,
+            req_info.items,
-                items_offset,
+            req_info.items_offset,
-                extend_prefix_len,
+            req_info.extend_prefix_len,
-                extend_seq_len,
+            req_info.extend_seq_len,
             input_ids,
             device,
         )
         if chunk_embedding is not None:
-                embedding_list.append(chunk_embedding)
+            all_chunks.append((req_info.req_idx, chunk_embedding))
+
+    # Sort by original request index to maintain correct output order
+    all_chunks.sort(key=lambda x: x[0])
+    embedding_list = [chunk for _, chunk in all_chunks]
 
     if len(embedding_list) == 0:
         return None, input_ids

python/sglang/srt/models/dots_vlm_vit.py

@@ -9,7 +9,11 @@ from torch.nn import LayerNorm
 from transformers.modeling_utils import PreTrainedModel
 
 from sglang.srt.configs.dots_vlm import DotsVisionConfig
-from sglang.srt.layers.attention.vision import VisionAttention
+from sglang.srt.layers.attention.vision import (
+    VisionAttention,
+    VisionAttentionMetadata,
+    prepare_vision_attention_metadata,
+)
 from sglang.srt.layers.conv import Conv2dLayer
 from sglang.srt.layers.quantization import QuantizationConfig
 from sglang.srt.runtime_context import get_parallel
@@ -173,11 +177,18 @@ class DotsVisionBlock(nn.Module):
         self.mlp = DotsSwiGLUFFN(config, quant_config)
         self.norm2 = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
 
-    def forward(self, hidden_states, cu_seqlens, rotary_pos_emb) -> torch.Tensor:
+    def forward(
+        self,
+        hidden_states,
+        cu_seqlens,
+        rotary_pos_emb,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
+    ) -> torch.Tensor:
         hidden_states = hidden_states + self.attn(
             self.norm1(hidden_states),
             cu_seqlens=cu_seqlens,
             position_embeddings=rotary_pos_emb,
+            forward_metadata=forward_metadata,
         )
         hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
         return hidden_states
@@ -319,10 +330,16 @@ class DotsVisionTransformer(PreTrainedModel):
         # cu_seqlens must be on cpu because of npu_flash_attention_unpad operator restriction
         if is_npu():
             cu_seqlens = cu_seqlens.to("cpu")
+        forward_metadata = prepare_vision_attention_metadata(
+            cu_seqlens, device=hidden_states.device
+        )
 
         for blk in self.blocks:
             hidden_states = blk(
-                hidden_states, cu_seqlens=cu_seqlens, rotary_pos_emb=rotary_pos_emb
+                hidden_states,
+                cu_seqlens=cu_seqlens,
+                rotary_pos_emb=rotary_pos_emb,
+                forward_metadata=forward_metadata,
             )
 
         if self.config.post_norm:

python/sglang/srt/models/ernie45_vl.py

@@ -24,7 +24,11 @@ from einops import rearrange
 from transformers import PretrainedConfig
 
 from sglang.srt.layers.activation import QuickGELU
-from sglang.srt.layers.attention.vision import VisionAttention
+from sglang.srt.layers.attention.vision import (
+    VisionAttention,
+    VisionAttentionMetadata,
+    prepare_vision_attention_metadata,
+)
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
 from sglang.srt.layers.logits_processor import LogitsProcessor
@@ -123,6 +127,7 @@ class Ernie4_5_VisionBlock(nn.Module):
         cu_seqlens: torch.Tensor,
         rotary_pos_emb_cos: torch.Tensor,
         rotary_pos_emb_sin: torch.Tensor,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
     ) -> torch.Tensor:
         hidden_states = self.norm1(x)
         hidden_states = rearrange(hidden_states, "s b ... -> b s ...")
@@ -131,6 +136,7 @@ class Ernie4_5_VisionBlock(nn.Module):
             cu_seqlens=cu_seqlens,
             rotary_pos_emb_cos=rotary_pos_emb_cos,
             rotary_pos_emb_sin=rotary_pos_emb_sin,
+            forward_metadata=forward_metadata,
         )
         attn = rearrange(attn, "b s ... -> s b ...")
         x = x + attn
@@ -481,6 +487,9 @@ class Ernie4_5_VisionTransformer(nn.Module):
             grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
         ).cumsum(dim=0, dtype=torch.int32)
         cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens])
+        forward_metadata = prepare_vision_attention_metadata(
+            cu_seqlens, device=x.device
+        )
 
         # transformers
         x = x.unsqueeze(1)
@@ -490,6 +499,7 @@ class Ernie4_5_VisionTransformer(nn.Module):
                 cu_seqlens=cu_seqlens,
                 rotary_pos_emb_cos=rotary_pos_emb_cos,
                 rotary_pos_emb_sin=rotary_pos_emb_sin,
+                forward_metadata=forward_metadata,
             )
 
         final_output = self.ln(x)

python/sglang/srt/models/glm4v.py

@@ -30,7 +30,11 @@ from transformers.models.glm4v.configuration_glm4v import Glm4vConfig, Glm4vVisi
 from sglang.srt.distributed.parallel_state import get_pp_group
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.attention import vision_utils
-from sglang.srt.layers.attention.vision import VisionAttention
+from sglang.srt.layers.attention.vision import (
+    VisionAttention,
+    VisionAttentionMetadata,
+    prepare_vision_attention_metadata,
+)
 from sglang.srt.layers.conv import Conv3dLayer
 from sglang.srt.layers.layernorm import LayerNorm, RMSNorm
 from sglang.srt.layers.linear import (
@@ -155,6 +159,7 @@ class Glm4vVisionBlock(nn.Module):
         cu_seqlens: torch.Tensor,
         rotary_pos_emb_cos: torch.Tensor,
         rotary_pos_emb_sin: torch.Tensor,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
     ) -> torch.Tensor:
         S, B, H = x.shape
         # norm1: flatten to 2D -> [S*B, H], then reshape back
@@ -168,6 +173,7 @@ class Glm4vVisionBlock(nn.Module):
             cu_seqlens=cu_seqlens,
             rotary_pos_emb_cos=rotary_pos_emb_cos,
             rotary_pos_emb_sin=rotary_pos_emb_sin,
+            forward_metadata=forward_metadata,
         )
         attn = rearrange(attn, "b s h -> s b h")
 
@@ -515,6 +521,10 @@ class Glm4vVisionModel(nn.Module):
         if is_npu():
             cu_seqlens = cu_seqlens.to("cpu")
 
+        forward_metadata = prepare_vision_attention_metadata(
+            cu_seqlens, device=self.device
+        )
+
         # x.shape: (s, b, d) where b=1 for vision processing
         # transformers
         x = x.unsqueeze(1)
@@ -524,6 +534,7 @@ class Glm4vVisionModel(nn.Module):
                 cu_seqlens=cu_seqlens,
                 rotary_pos_emb_cos=rotary_pos_emb_cos,
                 rotary_pos_emb_sin=rotary_pos_emb_sin,
+                forward_metadata=forward_metadata,
             )
 
         # adapter

python/sglang/srt/models/glm_image_vl.py

@@ -27,7 +27,11 @@ import torch.nn as nn
 import torch.nn.functional as F
 from einops import rearrange
 
-from sglang.srt.layers.attention.vision import VisionAttention
+from sglang.srt.layers.attention.vision import (
+    VisionAttention,
+    VisionAttentionMetadata,
+    prepare_vision_attention_metadata,
+)
 from sglang.srt.layers.dp_attention import is_dp_attention_enabled
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
@@ -196,11 +200,16 @@ class GlmImageVisionBlock(nn.Module):
         self,
         x: torch.Tensor,
         cu_seqlens: torch.Tensor,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
     ) -> torch.Tensor:
         # x shape: (S, B, H) where B=1
         hidden_states = self.norm1(x)
         hidden_states = rearrange(hidden_states, "s b ... -> b s ...")
-        attn = self.attn(hidden_states, cu_seqlens=cu_seqlens)
+        attn = self.attn(
+            hidden_states,
+            cu_seqlens=cu_seqlens,
+            forward_metadata=forward_metadata,
+        )
         attn = rearrange(attn, "b s ... -> s b ...")
         x = x + attn
 
@@ -294,6 +303,9 @@ class GlmImageVisionModel(nn.Module):
             cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)
         else:
             cu_seqlens = cu_seqlens.to("cpu")
+        forward_metadata = prepare_vision_attention_metadata(
+            cu_seqlens, device=hidden_states.device
+        )
 
         seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
 
@@ -309,7 +321,11 @@ class GlmImageVisionModel(nn.Module):
         hidden_states = hidden_states.unsqueeze(1)
 
         for blk in self.blocks:
-            hidden_states = blk(hidden_states, cu_seqlens=cu_seqlens)
+            hidden_states = blk(
+                hidden_states,
+                cu_seqlens=cu_seqlens,
+                forward_metadata=forward_metadata,
+            )
 
         # (S, 1, H) -> (S, H)
         return hidden_states.squeeze(1)

python/sglang/srt/models/glm_ocr.py

@@ -32,7 +32,11 @@ from transformers.models.glm_ocr.configuration_glm_ocr import (
 
 from sglang.srt.distributed.parallel_state import get_pp_group
 from sglang.srt.layers.attention import vision_utils
-from sglang.srt.layers.attention.vision import VisionAttention
+from sglang.srt.layers.attention.vision import (
+    VisionAttention,
+    VisionAttentionMetadata,
+    prepare_vision_attention_metadata,
+)
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.pooler import Pooler, PoolingType
@@ -112,6 +116,7 @@ class GlmOcrVisionBlock(nn.Module):
         cu_seqlens: torch.Tensor,
         rotary_pos_emb_cos: torch.Tensor,
         rotary_pos_emb_sin: torch.Tensor,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
     ) -> torch.Tensor:
         S, B, H = x.shape
         # norm1: flatten to 2D -> [S*B, H], then reshape back
@@ -125,6 +130,7 @@ class GlmOcrVisionBlock(nn.Module):
             cu_seqlens=cu_seqlens,
             rotary_pos_emb_cos=rotary_pos_emb_cos,
             rotary_pos_emb_sin=rotary_pos_emb_sin,
+            forward_metadata=forward_metadata,
         )
         attn = rearrange(attn, "b s h -> s b h")
 
@@ -236,6 +242,9 @@ class GlmOcrVisionModel(Glm4vVisionModel):
             grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
         ).cumsum(dim=0, dtype=torch.int32)
         cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens])
+        forward_metadata = prepare_vision_attention_metadata(
+            cu_seqlens, device=x.device
+        )
 
         rotary_pos_emb_cos = torch.cat([rotary_pos_emb_cos, rotary_pos_emb_cos], dim=-1)
         rotary_pos_emb_sin = torch.cat([rotary_pos_emb_sin, rotary_pos_emb_sin], dim=-1)
@@ -249,6 +258,7 @@ class GlmOcrVisionModel(Glm4vVisionModel):
                 cu_seqlens=cu_seqlens,
                 rotary_pos_emb_cos=rotary_pos_emb_cos,
                 rotary_pos_emb_sin=rotary_pos_emb_sin,
+                forward_metadata=forward_metadata,
             )
 
         # adapter

python/sglang/srt/models/idefics2.py

@@ -25,7 +25,11 @@ from torch import nn
 from transformers import PretrainedConfig
 
 from sglang.srt.layers.activation import get_act_fn
-from sglang.srt.layers.attention.vision import VisionAttention
+from sglang.srt.layers.attention.vision import (
+    VisionAttention,
+    VisionAttentionMetadata,
+    prepare_vision_attention_metadata,
+)
 from sglang.srt.layers.conv import Conv2dLayer
 from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
@@ -99,6 +103,7 @@ class Idefics2EncoderLayer(nn.Module):
         self,
         hidden_states: torch.Tensor,
         cu_seqlens: torch.Tensor,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
     ) -> torch.Tensor:
         """
         Args:
@@ -108,7 +113,11 @@ class Idefics2EncoderLayer(nn.Module):
         """
         residual = hidden_states
         hidden_states = self.layer_norm1(hidden_states)
-        hidden_states = self.self_attn(hidden_states, cu_seqlens=cu_seqlens)
+        hidden_states = self.self_attn(
+            hidden_states,
+            cu_seqlens=cu_seqlens,
+            forward_metadata=forward_metadata,
+        )
 
         hidden_states = residual + hidden_states
         residual = hidden_states
@@ -152,6 +161,7 @@ class Idefics2Encoder(nn.Module):
         self,
         inputs_embeds: torch.Tensor,
         cu_seqlens: torch.Tensor,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
     ) -> torch.Tensor:
         r"""
         Args:
@@ -163,13 +173,18 @@ class Idefics2Encoder(nn.Module):
                 internal embedding lookup matrix.
         """
         # cu_seqlens must be on cpu because of npu_flash_attention_unpad operator restriction
+        hidden_states = inputs_embeds
         if is_npu():
             cu_seqlens = cu_seqlens.to("cpu")
-        hidden_states = inputs_embeds
+        if forward_metadata is None:
+            forward_metadata = prepare_vision_attention_metadata(
+                cu_seqlens, device=hidden_states.device
+            )
         for encoder_layer in self.layers:
             layer_outputs = encoder_layer(
                 hidden_states,
                 cu_seqlens=cu_seqlens,
+                forward_metadata=forward_metadata,
             )
             hidden_states = layer_outputs
         return hidden_states
@@ -340,6 +355,9 @@ class Idefics2VisionTransformer(nn.Module):
         encoder_outputs = self.encoder(
             hidden_states,
             cu_seqlens=cu_seqlens,
+            forward_metadata=prepare_vision_attention_metadata(
+                cu_seqlens, device=hidden_states.device
+            ),
         )
         last_hidden_state = self.post_layernorm(encoder_outputs)
         return last_hidden_state

python/sglang/srt/models/internvl.py

@@ -14,7 +14,13 @@ from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPo
 from sglang.srt.environ import envs
 from sglang.srt.layers.activation import get_act_fn
 from sglang.srt.layers.attention import vision_utils
-from sglang.srt.layers.attention.vision import SingletonCache, VisionAttention
+from sglang.srt.layers.attention.vision import (
+    SingletonCache,
+    VisionAttention,
+    VisionAttentionMetadata,
+    _get_cu_seqlens_for_shape,
+    prepare_vision_attention_metadata,
+)
 from sglang.srt.layers.conv import Conv2dLayer
 from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
 from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
@@ -86,8 +92,14 @@ class InternAttention(nn.Module):
         hidden_states: torch.Tensor,
         cu_seqlens: torch.Tensor,
         output_ws: Optional[torch.Tensor] = None,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
     ) -> torch.Tensor:
-        out = self.attn(hidden_states, cu_seqlens=cu_seqlens, output_ws=output_ws)
+        out = self.attn(
+            hidden_states,
+            cu_seqlens=cu_seqlens,
+            output_ws=output_ws,
+            forward_metadata=forward_metadata,
+        )
         outs = self.proj_drop(out)
         return outs
 
@@ -259,11 +271,8 @@ class InternVisionEncoderLayer(nn.Module):
         hidden_states: torch.Tensor,
         cu_seqlens: torch.Tensor,
         output_ws: Optional[torch.Tensor] = None,
-    ) -> Tuple[
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
-        torch.FloatTensor,
+    ) -> torch.FloatTensor:
-        Optional[torch.FloatTensor],
-        Optional[Tuple[torch.FloatTensor]],
-    ]:
         """
         Args:
             hidden_states (`Tuple[torch.FloatTensor, Optional[torch.FloatTensor]]`): input to the layer of shape `(batch, seq_len, embed_dim)`
@@ -274,6 +283,7 @@ class InternVisionEncoderLayer(nn.Module):
                 self.norm1(hidden_states).to(hidden_states.dtype),
                 cu_seqlens=cu_seqlens,
                 output_ws=output_ws,
+                forward_metadata=forward_metadata,
             )
             * self.ls1
         )
@@ -344,7 +354,6 @@ class InternVisionEncoder(nn.Module):
                 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
         """
         if self.enable_cg and (not output_hidden_states):
-            # graph path only returns last_hidden_state
             hidden_states = inputs_embeds.to(device=inputs_embeds.device).contiguous()
             hidden_states = self.cuda_graph_runner.run(hidden_states)
             if not return_dict:
@@ -364,12 +373,26 @@ class InternVisionEncoder(nn.Module):
         hidden_states = inputs_embeds
 
         if cu_seqlens is None:
+            bsz, seq_len, _ = inputs_embeds.shape
+            cu_seqlens = _get_cu_seqlens_for_shape(
+                bsz, seq_len, device=inputs_embeds.device
+            )
+        forward_metadata = None
+        if isinstance(cu_seqlens, torch.Tensor):
+            forward_metadata = prepare_vision_attention_metadata(
+                cu_seqlens, device=inputs_embeds.device
+            )
+        elif cu_seqlens is None:
             cu_seqlens = SingletonCache()
 
         for idx, encoder_layer in enumerate(self.layers):
             if output_hidden_states:
                 encoder_states = encoder_states + (hidden_states,)
-            layer_outputs = encoder_layer(hidden_states, cu_seqlens=cu_seqlens)
+            layer_outputs = encoder_layer(
+                hidden_states,
+                cu_seqlens=cu_seqlens,
+                forward_metadata=forward_metadata,
+            )
             hidden_states = layer_outputs
 
         if output_hidden_states:

python/sglang/srt/models/kimi_k25.py

@@ -10,8 +10,13 @@ from transformers.activations import PytorchGELUTanh
 
 from sglang.srt.configs.kimi_k25 import KimiK25Config, KimiK25VisionConfig
 from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
-from sglang.srt.layers.attention.vision import VisionAttention
+from sglang.srt.layers.attention.vision import (
+    VisionAttention,
+    VisionAttentionMetadata,
+    prepare_vision_attention_metadata,
+)
 from sglang.srt.layers.conv import Conv2dLayer
+from sglang.srt.layers.dp_attention import is_dp_attention_enabled
 from sglang.srt.layers.linear import ReplicatedLinear
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig
@@ -39,8 +44,6 @@ from sglang.srt.utils import add_prefix, is_npu
 
 logger = logging.getLogger(__name__)
 
-from sglang.srt.layers.dp_attention import is_dp_attention_enabled
-
 _is_npu = is_npu()
 
 
@@ -145,6 +148,7 @@ class MoonViTEncoderLayer(nn.Module):
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
         rope_freqs_cis: torch.Tensor | None = None,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
         sequence_lengths: torch.Tensor | None = None,
     ):
         residual = hidden_states
@@ -154,6 +158,7 @@ class MoonViTEncoderLayer(nn.Module):
             hidden_states,
             cu_seqlens=cu_seqlens,
             position_embeddings=rope_freqs_cis,
+            forward_metadata=forward_metadata,
             max_seqlen=max_seqlen,
             sequence_lengths=sequence_lengths,
         )
@@ -484,12 +489,17 @@ class MoonViT3dEncoder(nn.Module):
         max_seqlen = int(lengths.max().item())
         cu_seqlens = lengths.to(hidden_states.device).cumsum(dim=0, dtype=torch.int32)
 
+        forward_metadata = prepare_vision_attention_metadata(
+            cu_seqlens, device=hidden_states.device
+        )
+
         for block in self.blocks:
             hidden_states = block(
                 hidden_states,
                 cu_seqlens,
                 max_seqlen,
                 rope_freqs_cis=rope_freqs_cis,
+                forward_metadata=forward_metadata,
                 sequence_lengths=sequence_lengths,
             )
 

python/sglang/srt/models/mimo_vl.py

@@ -14,7 +14,11 @@ from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
     Qwen2_5_VisionRotaryEmbedding,
 )
 
-from sglang.srt.layers.attention.vision import VisionAttention
+from sglang.srt.layers.attention.vision import (
+    VisionAttention,
+    VisionAttentionMetadata,
+    prepare_vision_attention_metadata,
+)
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.quantization import QuantizationConfig
 from sglang.srt.models.qwen2_5_vl import Qwen2_5_VisionPatchMerger, Qwen2_5_VLMLP
@@ -194,6 +198,7 @@ class MiMoVisionBlock(nn.Module):
         max_seqlen: int,
         position_embeddings: torch.Tensor,
         full_attn: bool = True,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
     ) -> torch.Tensor:
         S, B, H = x.shape
         # norm1: flatten to 2D -> [S*B, H], then reshape back
@@ -208,6 +213,7 @@ class MiMoVisionBlock(nn.Module):
             max_seqlen=max_seqlen,
             position_embeddings=position_embeddings,
             full_attn=full_attn,
+            forward_metadata=forward_metadata,
         )
         attn = rearrange(attn, "b s h -> s b h")
 
@@ -429,6 +435,9 @@ class MiMoVisionTransformer(nn.Module):
             ]
         )
         max_seqlen = seqlens.max().item()
+        forward_metadata = prepare_vision_attention_metadata(
+            cu_seqlens, device=x.device
+        )
 
         row_based_embeddings = get_position_embeddings(emb, x)
         col_based_embeddings = get_position_embeddings(
@@ -446,6 +455,7 @@ class MiMoVisionTransformer(nn.Module):
             reverse_window_index_1d_col,
             cu_seqlens,
             max_seqlen,
+            forward_metadata,
         )
 
     def run_blocks(
@@ -457,6 +467,7 @@ class MiMoVisionTransformer(nn.Module):
         reverse_window_index_1d_col: torch.Tensor,
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
+        forward_metadata: VisionAttentionMetadata,
     ) -> torch.Tensor:
         for layer_num, blk in enumerate(self.blocks):
             window_attn_type = self.vit_window_attn_types[layer_num]
@@ -485,6 +496,7 @@ class MiMoVisionTransformer(nn.Module):
                 max_seqlen=max_seqlen,
                 position_embeddings=position_embeddings,
                 full_attn=full_attn,
+                forward_metadata=forward_metadata,
             )
         x = self.merger(x)
         return x
@@ -502,6 +514,7 @@ class MiMoVisionTransformer(nn.Module):
             reverse_window_index_1d_col,
             cu_seqlens,
             max_seqlen,
+            forward_metadata,
         ) = self._prepare_forward(x, grid_thw)
 
         return self.run_blocks(
@@ -512,4 +525,5 @@ class MiMoVisionTransformer(nn.Module):
             reverse_window_index_1d_col,
             cu_seqlens,
             max_seqlen,
+            forward_metadata,
         )

python/sglang/srt/models/minicpmv_vit.py

@@ -31,7 +31,10 @@ from torch import nn
 from transformers import PretrainedConfig
 
 from sglang.srt.layers.activation import get_act_fn
-from sglang.srt.layers.attention.vision import VisionAttention
+from sglang.srt.layers.attention.vision import (
+    VisionAttention,
+    prepare_vision_attention_metadata,
+)
 from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.models.idefics2 import (
@@ -170,7 +173,14 @@ class MiniCPMV_ViTWindowAttentionMerger(nn.Module):
             window_cu_seqlens = window_cu_seqlens.to("cpu")
 
         hidden_states = hidden_states[:, window_index, :]
-        hidden_states = self.self_attn(hidden_states, cu_seqlens=window_cu_seqlens)
+        window_metadata = prepare_vision_attention_metadata(
+            window_cu_seqlens, device=hidden_states.device
+        )
+        hidden_states = self.self_attn(
+            hidden_states,
+            cu_seqlens=window_cu_seqlens,
+            forward_metadata=window_metadata,
+        )
         hidden_states = hidden_states[:, torch.argsort(window_index), :]
         hidden_states = residual + hidden_states
 
@@ -501,13 +511,20 @@ class MiniCPMV_VisionTransformer(nn.Module):
         cu_seqlens, max_seqlens = self.compute_cu_seqlens(target_sizes)
         if is_npu():
             cu_seqlens = cu_seqlens.to("cpu")
+        forward_metadata = prepare_vision_attention_metadata(
+            cu_seqlens, device=hidden_states.device
+        )
 
         if use_vit_merger:
             # Encoder loop lives here (not inside ``MiniCPMV_VisionEncoder``)
             # so we can fire ``vit_merger`` after layer ``insert_layer_id``
             # without coupling the encoder module to it.
             for layer_index, layer in enumerate(self.encoder.layers):
-                hidden_states = layer(hidden_states, cu_seqlens=cu_seqlens)
+                hidden_states = layer(
+                    hidden_states,
+                    cu_seqlens=cu_seqlens,
+                    forward_metadata=forward_metadata,
+                )
                 if layer_index == self.insert_layer_id:
                     (
                         hidden_states,
@@ -519,8 +536,15 @@ class MiniCPMV_VisionTransformer(nn.Module):
                     )
                     if is_npu():
                         cu_seqlens = cu_seqlens.to("cpu")
+                    forward_metadata = prepare_vision_attention_metadata(
+                        cu_seqlens, device=hidden_states.device
+                    )
         else:
-            hidden_states = self.encoder(hidden_states, cu_seqlens=cu_seqlens)
+            hidden_states = self.encoder(
+                hidden_states,
+                cu_seqlens=cu_seqlens,
+                forward_metadata=forward_metadata,
+            )
 
         hidden_states = self.post_layernorm(hidden_states)
         return hidden_states, target_sizes

python/sglang/srt/models/minimax_vl_common.py

@@ -14,6 +14,8 @@ from sglang.srt.layers.attention.vision import (
     FLASHINFER_MAX_SEQLEN_BUCKETS,
     FLASHINFER_WORKSPACE_SIZE_BYTES,
     VisionAttention,
+    VisionAttentionMetadata,
+    prepare_vision_attention_metadata,
 )
 from sglang.srt.layers.dp_attention import is_dp_attention_enabled
 from sglang.srt.layers.linear import (
@@ -305,6 +307,7 @@ class CLIPEncoderLayer(nn.Module):
         rotary_pos_emb: torch.Tensor,
         max_seqlen: Optional[int] = None,
         sequence_lengths: Optional[torch.Tensor] = None,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
     ) -> torch.Tensor:
         residual = hidden_states
         hidden_states = self.layer_norm1(hidden_states)
@@ -314,6 +317,7 @@ class CLIPEncoderLayer(nn.Module):
             position_embeddings=rotary_pos_emb,
             max_seqlen=max_seqlen,
             sequence_lengths=sequence_lengths,
+            forward_metadata=forward_metadata,
         )
         hidden_states = residual + hidden_states
 
@@ -361,6 +365,7 @@ class CLIPEncoder(nn.Module):
         rotary_pos_emb: torch.Tensor,
         max_seqlen: Optional[int] = None,
         sequence_lengths: Optional[torch.Tensor] = None,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
     ) -> torch.Tensor:
         hidden_states = inputs_embeds
         cos_sin = _prepare_rotary_cos_sin(rotary_pos_emb)
@@ -372,6 +377,7 @@ class CLIPEncoder(nn.Module):
                 cos_sin,
                 max_seqlen=max_seqlen,
                 sequence_lengths=sequence_lengths,
+                forward_metadata=forward_metadata,
             )
 
         return hidden_states
@@ -680,12 +686,25 @@ class MiniMaxVLVisionTransformer(nn.Module):
                 max_seqlen,
             ) = self._build_flashinfer_cudnn_inputs(cu_seq_len)
 
+        forward_metadata = prepare_vision_attention_metadata(
+            cu_seq_len,
+            device=hidden_states.device,
+            packed_indptrs=(
+                encoder_cu_seq_len
+                if get_server_args().mm_attention_backend == "flashinfer_cudnn"
+                else None
+            ),
+            sequence_lengths=sequence_lengths,
+            flashinfer_max_seqlen=max_seqlen,
+        )
+
         return self.encoder(
             inputs_embeds=hidden_states,
             cu_seq_len=encoder_cu_seq_len,
             rotary_pos_emb=rotary_pos_emb,
             max_seqlen=max_seqlen,
             sequence_lengths=sequence_lengths,
+            forward_metadata=forward_metadata,
         )
 
 

python/sglang/srt/models/moss_vl.py

@@ -17,7 +17,11 @@ from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
 
 from sglang.srt.environ import envs
 from sglang.srt.layers.activation import SiluAndMul
-from sglang.srt.layers.attention.vision import VisionAttention
+from sglang.srt.layers.attention.vision import (
+    VisionAttention,
+    VisionAttentionMetadata,
+    prepare_vision_attention_metadata,
+)
 from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes
 from sglang.srt.layers.conv import Conv3dLayer
 from sglang.srt.layers.layernorm import RMSNorm
@@ -163,6 +167,7 @@ class MossVLVisionBlock(nn.Module):
         x: torch.Tensor,
         cu_seqlens: torch.Tensor,
         position_embeddings: torch.Tensor,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
     ) -> torch.Tensor:
         hidden_states = self.norm1(x)
         hidden_states = rearrange(hidden_states, "s b ... -> b s ...")
@@ -170,6 +175,7 @@ class MossVLVisionBlock(nn.Module):
             hidden_states,
             cu_seqlens=cu_seqlens,
             position_embeddings=position_embeddings,
+            forward_metadata=forward_metadata,
         )
         attn = rearrange(attn, "b s ... -> s b ...")
         x = x + attn
@@ -545,12 +551,20 @@ class MossVLVisionModel(nn.Module):
                 cu_seqlens.to(torch.int32),
             ]
         )
+        forward_metadata = prepare_vision_attention_metadata(
+            cu_seqlens, device=x.device
+        )
 
         x = x.unsqueeze(1)
 
         deepstack_features = []
         for layer_idx, blk in enumerate(self.blocks):
-            x = blk(x, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings)
+            x = blk(
+                x,
+                cu_seqlens=cu_seqlens,
+                position_embeddings=position_embeddings,
+                forward_metadata=forward_metadata,
+            )
             if layer_idx in self.deepstack_visual_indexes:
                 deepstack_features.append(x)
 

python/sglang/srt/models/paddleocr_vl.py

@@ -25,7 +25,11 @@ from transformers.activations import GELUActivation
 from transformers.utils import torch_int
 
 from sglang.srt.layers.activation import get_act_fn
-from sglang.srt.layers.attention.vision import VisionAttention
+from sglang.srt.layers.attention.vision import (
+    VisionAttention,
+    VisionAttentionMetadata,
+    prepare_vision_attention_metadata,
+)
 from sglang.srt.layers.conv import Conv2dLayer
 from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
@@ -345,6 +349,7 @@ class SiglipEncoderLayer(nn.Module):
         hidden_states: torch.Tensor,
         cu_seqlens: Optional[List[torch.Tensor]] = None,
         rope_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
     ) -> Tuple[torch.FloatTensor]:
 
         residual = hidden_states
@@ -355,6 +360,7 @@ class SiglipEncoderLayer(nn.Module):
             hidden_states,
             cu_seqlens=cu_seqlens,
             position_embeddings=rope_emb,
+            forward_metadata=forward_metadata,
         )
 
         hidden_states = residual + hidden_states
@@ -447,6 +453,9 @@ class SiglipEncoder(nn.Module):
         if is_npu() and isinstance(cu_seqlens, torch.Tensor):
             cu_seqlens = cu_seqlens.to("cpu")
         attn_cu_seqlens = cu_seqlens
+        forward_metadata = prepare_vision_attention_metadata(
+            attn_cu_seqlens, device=hidden_states.device
+        )
         hidden_states = inputs_embeds
 
         for encoder_layer in self.layers:
@@ -454,6 +463,7 @@ class SiglipEncoder(nn.Module):
                 hidden_states,
                 cu_seqlens=attn_cu_seqlens,
                 rope_emb=rope_emb,
+                forward_metadata=forward_metadata,
             )
         return hidden_states
 

python/sglang/srt/models/qwen2_5_vl.py

@@ -45,7 +45,11 @@ from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
 from sglang.srt.distributed.parallel_state import get_pp_group
 from sglang.srt.environ import envs
 from sglang.srt.layers.activation import SiluAndMul
-from sglang.srt.layers.attention.vision import VisionAttention
+from sglang.srt.layers.attention.vision import (
+    VisionAttention,
+    VisionAttentionMetadata,
+    prepare_vision_attention_metadata,
+)
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
     ColumnParallelLinear,
@@ -180,6 +184,7 @@ class Qwen2_5_VisionBlock(nn.Module):
         cu_seqlens: torch.Tensor,
         position_embeddings: torch.Tensor,
         output_ws=None,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
     ) -> torch.Tensor:
         S, B, H = x.shape
         # norm1: flatten to 2D -> [S*B, H], then reshape back
@@ -193,6 +198,7 @@ class Qwen2_5_VisionBlock(nn.Module):
             cu_seqlens=cu_seqlens,
             position_embeddings=position_embeddings,
             output_ws=output_ws,
+            forward_metadata=forward_metadata,
         )
         attn = rearrange(attn, "b s h -> s b h")
 
@@ -469,6 +475,13 @@ class Qwen2_5_VisionTransformer(nn.Module, RotaryPosMixin):
         if is_npu():
             cu_seqlens = cu_seqlens.to("cpu")
             cu_window_seqlens = cu_window_seqlens.to("cpu")
+
+        # pre-compute attention metadata once for all layers (two variants)
+        full_metadata = prepare_vision_attention_metadata(cu_seqlens, device=x.device)
+        window_metadata = prepare_vision_attention_metadata(
+            cu_window_seqlens, device=x.device
+        )
+
         # transformers
         x = x.unsqueeze(1)
         for layer_num, blk in enumerate(self.blocks):
@@ -477,10 +490,15 @@ class Qwen2_5_VisionTransformer(nn.Module, RotaryPosMixin):
                 fullatt_indexes = fullatt_indexes.tolist()
             if layer_num in fullatt_indexes:
                 cu_seqlens_now = cu_seqlens
+                metadata_now = full_metadata
             else:
                 cu_seqlens_now = cu_window_seqlens
+                metadata_now = window_metadata
             x = blk(
-                x, cu_seqlens=cu_seqlens_now, position_embeddings=position_embeddings
+                x,
+                cu_seqlens=cu_seqlens_now,
+                position_embeddings=position_embeddings,
+                forward_metadata=metadata_now,
             )
 
         # adapter
@@ -518,8 +536,8 @@ class Qwen2_5_VisionTransformer(nn.Module, RotaryPosMixin):
 
         # [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[window_index, :, :]
-        x = x.reshape(seq_len, -1)  # [seq_len, hidden]
+        x = x.reshape(seq_len, -1)
 
         rotary_pos_emb = rotary_pos_emb.reshape(
             seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1
@@ -529,8 +547,6 @@ class Qwen2_5_VisionTransformer(nn.Module, RotaryPosMixin):
 
         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),

python/sglang/srt/models/qwen2_vl.py

@@ -34,7 +34,11 @@ from transformers import Qwen2VLConfig
 from transformers.models.qwen2_vl.configuration_qwen2_vl import Qwen2VLVisionConfig
 
 from sglang.srt.layers.activation import QuickGELU
-from sglang.srt.layers.attention.vision import VisionAttention
+from sglang.srt.layers.attention.vision import (
+    VisionAttention,
+    VisionAttentionMetadata,
+    prepare_vision_attention_metadata,
+)
 from sglang.srt.layers.conv import Conv3dLayer
 from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
 from sglang.srt.layers.logits_processor import LogitsProcessor
@@ -162,6 +166,7 @@ class Qwen2VisionBlock(nn.Module):
         x: torch.Tensor,
         cu_seqlens: torch.Tensor,
         position_embeddings: torch.Tensor,
+        forward_metadata: Optional["VisionAttentionMetadata"] = None,
     ) -> torch.Tensor:
         hidden_states = self.norm1(x)
         hidden_states = rearrange(hidden_states, "s b ... -> b s ...")
@@ -169,6 +174,7 @@ class Qwen2VisionBlock(nn.Module):
             hidden_states,
             cu_seqlens=cu_seqlens,
             position_embeddings=position_embeddings,
+            forward_metadata=forward_metadata,
         )
         attn = rearrange(attn, "b s ... -> s b ...")
         x = x + attn
@@ -398,10 +404,20 @@ class Qwen2VisionTransformer(nn.Module):
         if is_npu():
             cu_seqlens = cu_seqlens.to("cpu")
 
+        # pre-compute attention metadata once for all layers
+        forward_metadata = prepare_vision_attention_metadata(
+            cu_seqlens, device=self.device
+        )
+
         # transformers
         x = x.unsqueeze(1)
         for blk in self.blocks:
-            x = blk(x, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings)
+            x = blk(
+                x,
+                cu_seqlens=cu_seqlens,
+                position_embeddings=position_embeddings,
+                forward_metadata=forward_metadata,
+            )
 
         # adapter
         x = self.merger(x)

python/sglang/srt/models/qwen3_vl.py

@@ -34,6 +34,8 @@ from sglang.srt.layers.attention.vision import (
     FLASHINFER_MAX_SEQLEN_BUCKETS,
     FLASHINFER_WORKSPACE_SIZE_BYTES,
     VisionAttention,
+    VisionAttentionMetadata,
+    prepare_vision_attention_metadata,
 )
 from sglang.srt.layers.conv import Conv3dLayer
 from sglang.srt.layers.dp_attention import (
@@ -87,7 +89,6 @@ if _is_npu:
 
     graph_runners_dict["npu"] = ViTNpuGraphRunner
 
-
 logger = logging.getLogger(__name__)
 
 _is_cpu_amx_available = cpu_has_amx_support()
@@ -224,6 +225,7 @@ class Qwen3_VisionBlock(nn.Module):
         rotary_pos_emb_cos: torch.Tensor,
         rotary_pos_emb_sin: torch.Tensor,
         output_ws: Optional[torch.Tensor] = None,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
         max_seqlen: Optional[torch.Tensor] = None,
         sequence_lengths: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
@@ -235,6 +237,7 @@ 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,
+            forward_metadata=forward_metadata,
             max_seqlen=max_seqlen,
             sequence_lengths=sequence_lengths,
         )
@@ -909,8 +912,11 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
             [np.zeros(1, dtype=np.int32), token_cu_seqlens]
         )
 
+        # ---- pre-compute attention metadata once for all layers ----
+        packed_indptrs = None
+        flashinfer_sequence_lengths = None
         flashinfer_max_seqlen = 0
-        cu_seqlens = None
+
         if get_server_args().mm_attention_backend == "flashinfer_cudnn":
             # real token lens (B,)
             real_seq_lens = token_cu_seqlens[1:] - token_cu_seqlens[:-1]
@@ -924,11 +930,8 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
             )
 
             # 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 = (
+            elem_per_token = self.hidden_size // attn_tp_size
-                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(
@@ -936,31 +939,31 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
                 elem_per_token=elem_per_token,
             )
 
-            sequence_lengths = (
+            flashinfer_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
+            )
-
+            packed_indptrs = torch.from_numpy(offsets_packed).to(
-            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
+
+        forward_metadata = prepare_vision_attention_metadata(
+            cu_seqlens,
+            device=self.device,
+            packed_indptrs=packed_indptrs,
+            sequence_lengths=flashinfer_sequence_lengths,
+            flashinfer_max_seqlen=flashinfer_max_seqlen,
+        )
 
         x = x.unsqueeze(1)
 
-        cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)
-
         deepstack_feature_lists = []
         num_deepstack_captured = 0
 
@@ -970,8 +973,7 @@ 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,
+                forward_metadata=forward_metadata,
-                sequence_lengths=sequence_lengths,
             )
 
             if layer_num in self.deepstack_visual_indexes:

python/sglang/srt/models/siglip2.py

@@ -30,7 +30,11 @@ import torch.nn.functional as F
 from transformers import Siglip2VisionConfig
 
 from sglang.srt.layers.activation import get_act_fn
-from sglang.srt.layers.attention.vision import VisionAttention
+from sglang.srt.layers.attention.vision import (
+    VisionAttention,
+    VisionAttentionMetadata,
+    prepare_vision_attention_metadata,
+)
 from sglang.srt.layers.linear import (
     ColumnParallelLinear,
     RowParallelLinear,
@@ -203,6 +207,7 @@ class Siglip2Attention(nn.Module):
         hidden_states: torch.Tensor,
         cu_seqlens: torch.Tensor,
         max_seqlen: int | torch.Tensor,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
     ) -> torch.Tensor:
         """Forward pass with variable-length attention.
 
@@ -214,7 +219,11 @@ class Siglip2Attention(nn.Module):
         Returns:
             (1, total_tokens, embed_dim) attention output
         """
-        return self.attn(hidden_states, cu_seqlens=cu_seqlens)
+        return self.attn(
+            hidden_states,
+            cu_seqlens=cu_seqlens,
+            forward_metadata=forward_metadata,
+        )
 
 
 class Siglip2MLP(nn.Module):
@@ -279,6 +288,7 @@ class Siglip2EncoderLayer(nn.Module):
         hidden_states: torch.Tensor,
         cu_seqlens: torch.Tensor,
         max_seqlen: int | torch.Tensor,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
     ) -> torch.Tensor:
         """Forward pass for encoder layer.
 
@@ -294,6 +304,7 @@ class Siglip2EncoderLayer(nn.Module):
             hidden_states=hidden_states,
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
+            forward_metadata=forward_metadata,
         )
         hidden_states = residual + hidden_states
 
@@ -339,15 +350,21 @@ class Siglip2Encoder(nn.Module):
         cu_seqlens: torch.Tensor,
         max_seqlen: int | torch.Tensor,
         return_all_hidden_states: bool = False,
+        forward_metadata: Optional[VisionAttentionMetadata] = None,
     ) -> torch.Tensor | list[torch.Tensor]:
         hidden_states_pool = [inputs_embeds]
         hidden_states = inputs_embeds
+        if forward_metadata is None:
+            forward_metadata = prepare_vision_attention_metadata(
+                cu_seqlens, device=hidden_states.device
+            )
 
         for encoder_layer in self.layers:
             hidden_states = encoder_layer(
                 hidden_states,
                 cu_seqlens=cu_seqlens,
                 max_seqlen=max_seqlen,
+                forward_metadata=forward_metadata,
             )
             if return_all_hidden_states:
                 hidden_states_pool.append(hidden_states)
@@ -464,12 +481,16 @@ class Siglip2VisionTransformer(nn.Module):
             Vision features tensor
         """
         hidden_states = self.embeddings(pixel_values_packed, spatial_shapes)
+        forward_metadata = prepare_vision_attention_metadata(
+            cu_seqlens, device=hidden_states.device
+        )
 
         encoder_outputs = self.encoder(
             inputs_embeds=hidden_states,
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
             return_all_hidden_states=select_layers is not None,
+            forward_metadata=forward_metadata,
         )
 
         encoder_outputs = resolve_visual_encoder_outputs(

python/sglang/srt/multimodal/vit_cuda_graph_runner.py

@@ -153,6 +153,25 @@ class ViTCudaGraphRunner:
 
         override_backend = get_server_args().mm_attention_backend
 
+        if self._fullatt_block_indexes and 0 not in vit.fullatt_block_indexes:
+            warmup_cu_ws = [cu_window, cu_window_kk, max_window_len]
+        else:
+            warmup_cu_ws = [cu_full, cu_full_kk, max_full_len]
+        if override_backend == "fa3":
+            warmup_cu_ws = [warmup_cu_ws[0], warmup_cu_ws[2]]
+
+        warmup_kwargs = dict(
+            cu_seqlens=warmup_cu_ws, output_ws=self.block_ws[graph_key]
+        )
+        if position_embeddings is not None:
+            warmup_kwargs["position_embeddings"] = position_embeddings
+        elif rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None:
+            warmup_kwargs["rotary_pos_emb_cos"] = rotary_pos_emb_cos
+            warmup_kwargs["rotary_pos_emb_sin"] = rotary_pos_emb_sin
+        with torch.no_grad():
+            vit.blocks[0](self.block_input[graph_key], **warmup_kwargs)
+        torch.cuda.synchronize()
+
         with self._capture_context(), torch.cuda.graph(graph):
             y = None
             deepstack_outs: List[torch.Tensor] = []
@@ -290,6 +309,8 @@ class ViTCudaGraphRunner:
                 self.cu_full_len[graph_key] = cu_seqlens
                 self.cu_full_len_kk[graph_key] = cu_seqlens[1:] - cu_seqlens[:-1]
 
+        self.block_input[graph_key].copy_(x_3d)
+
         if position_embeddings is not None:
             # make sure rotary workspace
             head_dim = position_embeddings[0].shape[1]

test/registered/unit/layers/attention/test_vision_max_seqlen.py

@@ -25,6 +25,7 @@ def test_vision_flash3_uses_precomputed_max_seqlen(monkeypatch):
     monkeypatch.setattr(
         vision, "flash_attn_varlen_func", fake_flash_attn, raising=False
     )
+    monkeypatch.setattr(vision, "flash_attn_func", fake_flash_attn, raising=False)
 
     attention = vision.VisionFlash3Attention(use_data_parallel=True)
     q = torch.zeros(3, 1, 8)
@@ -89,6 +90,7 @@ def test_vision_flash4_uses_precomputed_max_seqlen(monkeypatch):
     monkeypatch.setattr(
         vision, "flash_attn_varlen_func", fake_flash_attn, raising=False
     )
+    monkeypatch.setattr(vision, "flash_attn_func", fake_flash_attn, raising=False)
 
     attention = vision.VisionFlash4Attention(use_data_parallel=True)
     q = torch.zeros(3, 1, 8)
@@ -154,6 +156,7 @@ def test_kimi_moonvit_precomputes_sequence_lengths_once():
             max_seqlen,
             rope_freqs_cis,
             sequence_lengths,
+            forward_metadata=None,
         ):
             recorded["cu_seqlens"] = cu_seqlens
             recorded["max_seqlen"] = max_seqlen