vllm.model_executor.models.llava_onevision2 ¶

@MULTIMODAL_REGISTRY.register_processor(
    LlavaOnevision2MultiModalProcessor,
    info=LlavaOnevision2ProcessingInfo,
    dummy_inputs=LlavaOnevision2DummyInputsBuilder,
)
class LlavaOnevision2ForConditionalGeneration(
    nn.Module, SupportsMultiModal, SupportsPP
):
    """vLLM-side OV2 top-level model.

    Weight name rewriting (HF checkpoint → vLLM module tree):
      Prefix rewrites only (longest match wins). Vision tower attribute names
      mirror HF names verbatim, so no substring rules are needed. Substring
      rules would otherwise collide with the Qwen3 text-path ``self_attn``
      modules and break the language-model loader.
        ``model.language_model.``        → ``language_model.model.``
        ``model.visual.``                → ``visual.``
        ``lm_head.``                     → ``language_model.lm_head.``
        ``model.`` (fallback)            → ``language_model.model.``
    """

    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_prefix={
            "model.language_model.": "language_model.model.",
            "model.visual.": "visual.",
            "lm_head.": "language_model.lm_head.",
            "model.": "language_model.model.",
        }
    )

    @classmethod
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
        if modality.startswith("image"):
            return "<|vision_start|><|image_pad|><|vision_end|>"
        if modality.startswith("video"):
            return "<|vision_start|><|video_pad|><|vision_end|>"
        raise ValueError("Only image or video modality is supported")

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
        super().__init__()
        config = vllm_config.model_config.hf_config
        multimodal_config = vllm_config.model_config.multimodal_config

        self.config = config
        self.multimodal_config = multimodal_config

        # Build the vision tower under the tower marker so it is shared by the
        # image, video-frame and video-codec backends (all routed through
        # ``self.visual``). When both modalities are disabled via
        # ``--limit-mm-per-prompt`` the marker turns the tower into a skipped
        # placeholder whose weights are dropped automatically by the loader.
        with self._mark_tower_model(vllm_config, {"image", "video"}):
            self.visual = LlavaOnevision2VisionTower(
                config.vision_config,
                text_hidden_size=config.text_config.hidden_size,
                norm_eps=getattr(config.vision_config, "layer_norm_eps", 1e-6),
                quant_config=None,
                prefix=maybe_prefix(prefix, "visual"),
            )

        # OV2 LLM is plain Qwen3 -- 1-D positions, no M-RoPE. The wrapper
        # LlavaOnevision2Config keeps text/vision configs nested (unlike OV1.5
        # which promoted text fields), so explicitly hand text_config to the
        # Qwen3 init path or qwen3.py will hit AttributeError on vocab_size.
        self.language_model = init_vllm_registered_model(
            vllm_config=vllm_config,
            hf_config=config.text_config,
            prefix=maybe_prefix(prefix, "language_model"),
            architectures=["Qwen3ForCausalLM"],
        )

        self.make_empty_intermediate_tensors = (
            self.language_model.make_empty_intermediate_tensors
        )

    def _validate_and_reshape_mm_tensor(
        self, mm_input: object, name: str
    ) -> torch.Tensor:
        if not isinstance(mm_input, (torch.Tensor, list)):
            raise ValueError(f"Incorrect type of {name}: {type(mm_input)}")
        if isinstance(mm_input, torch.Tensor):
            if mm_input.ndim == 2:
                return mm_input
            if mm_input.ndim != 3:
                raise ValueError(
                    f"{name} must be 2D or batched-3D, got shape={mm_input.shape}"
                )
            # Flatten the leading batch dim into the patch dim: (b, n, d) ->
            # (b*n, d), avoiding a Python list of row tensors.
            return mm_input.flatten(0, 1)
        return torch.concat(mm_input)

    def _parse_and_validate_image_input(
        self, **kwargs: object
    ) -> LlavaOnevision2ImageInputs | None:
        pixel_values = kwargs.pop("pixel_values", None)
        image_embeds = kwargs.pop("image_embeds", None)
        image_grid_thw = kwargs.pop("image_grid_thw", None)
        patch_positions = kwargs.pop("patch_positions", None)

        if pixel_values is None and image_embeds is None:
            return None

        if pixel_values is not None:
            pixel_values = self._validate_and_reshape_mm_tensor(
                pixel_values, "image pixel values"
            )
            image_grid_thw = self._validate_and_reshape_mm_tensor(
                image_grid_thw, "image grid_thw"
            )
            if patch_positions is None:
                raise ValueError(
                    "OV2 requires patch_positions alongside pixel_values; "
                    "ensure the HF processor produces it (image, video-"
                    "frames, and video-codec backends all do)."
                )
            patch_positions = self._validate_and_reshape_mm_tensor(
                patch_positions, "image patch_positions"
            )
            return LlavaOnevision2ImagePixelInputs(
                type="pixel_values",
                pixel_values=pixel_values,
                image_grid_thw=image_grid_thw,
                patch_positions=patch_positions,
            )

        # image_embeds path
        image_embeds = self._validate_and_reshape_mm_tensor(
            image_embeds, "image embeds"
        )
        image_grid_thw = self._validate_and_reshape_mm_tensor(
            image_grid_thw, "image grid_thw"
        )
        return LlavaOnevision2ImageEmbeddingInputs(
            type="image_embeds",
            image_embeds=image_embeds,
            image_grid_thw=image_grid_thw,
        )

    def _process_image_input(
        self, image_input: LlavaOnevision2ImageInputs
    ) -> tuple[torch.Tensor, ...]:
        grid_thw = image_input["image_grid_thw"]
        assert grid_thw.ndim == 2

        if image_input["type"] == "image_embeds":
            image_embeds = image_input["image_embeds"]
        else:
            image_embeds = self.visual(
                image_input["pixel_values"],
                grid_thw=grid_thw,
                patch_positions=image_input["patch_positions"],
            )

        merge_size = self.visual.spatial_merge_size
        sizes = grid_thw.prod(-1) // merge_size // merge_size
        return image_embeds.split(sizes.tolist())

    def _parse_and_validate_video_input(
        self, **kwargs: object
    ) -> LlavaOnevision2VideoPixelInputs | None:
        pixel_values_videos = kwargs.pop("pixel_values_videos", None)
        video_grid_thw = kwargs.pop("video_grid_thw", None)
        patch_positions_videos = kwargs.pop("patch_positions_videos", None)
        video_num_frames = kwargs.pop("video_num_frames", None)
        kwargs.pop("frame_timestamps", None)
        kwargs.pop("video_is_codec", None)

        if pixel_values_videos is None:
            return None

        pixel_values_videos = self._validate_and_reshape_mm_tensor(
            pixel_values_videos, "video pixel values"
        )
        video_grid_thw = self._validate_and_reshape_mm_tensor(
            video_grid_thw, "video grid_thw"
        )
        if patch_positions_videos is None:
            raise ValueError(
                "OV2 requires patch_positions_videos alongside pixel_values_videos."
            )
        patch_positions_videos = self._validate_and_reshape_mm_tensor(
            patch_positions_videos, "video patch_positions"
        )
        if video_num_frames is None:
            raise ValueError(
                "OV2 requires video_num_frames alongside pixel_values_videos."
            )
        if isinstance(video_num_frames, list):
            video_num_frames = torch.cat([v.flatten() for v in video_num_frames])
        else:
            video_num_frames = video_num_frames.flatten()
        return LlavaOnevision2VideoPixelInputs(
            type="pixel_values_videos",
            pixel_values_videos=pixel_values_videos,
            video_grid_thw=video_grid_thw,
            patch_positions_videos=patch_positions_videos,
            video_num_frames=video_num_frames,
        )

    def _process_video_input(
        self,
        video_input: LlavaOnevision2VideoPixelInputs,
    ) -> tuple[torch.Tensor, ...]:
        # OV2 encodes each video frame independently through the same
        # vision stack as still images, then splits the resulting token
        # stream into per-video chunks using video_num_frames.
        grid_thw = video_input["video_grid_thw"]
        assert grid_thw.ndim == 2
        video_embeds = self.visual(
            video_input["pixel_values_videos"],
            grid_thw=grid_thw,
            patch_positions=video_input["patch_positions_videos"],
        )
        merge_size = self.visual.spatial_merge_size
        per_frame_tokens = grid_thw.prod(-1) // merge_size // merge_size
        # Aggregate per-frame token counts into per-video token counts.
        num_frames = video_input["video_num_frames"].tolist()
        sizes: list[int] = []
        cursor = 0
        for n in num_frames:
            n = int(n)
            sizes.append(int(per_frame_tokens[cursor : cursor + n].sum()))
            cursor += n
        return video_embeds.split(sizes)

    def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
        modalities = {}
        for key in kwargs:
            if key in ("pixel_values", "image_embeds") and "images" not in modalities:
                modalities["images"] = self._parse_and_validate_image_input(**kwargs)
            if key == "pixel_values_videos" and "videos" not in modalities:
                modalities["videos"] = self._parse_and_validate_video_input(**kwargs)
        return modalities

    def get_language_model(self) -> torch.nn.Module:
        return self.language_model

    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
        modalities = self._parse_and_validate_multimodal_inputs(**kwargs)
        if not modalities:
            return []
        multimodal_embeddings: tuple[torch.Tensor, ...] = ()
        for modality in modalities:
            if modality == "images":
                multimodal_embeddings += self._process_image_input(modalities["images"])
            elif modality == "videos":
                multimodal_embeddings += self._process_video_input(modalities["videos"])
        return multimodal_embeddings

    def get_input_embeddings(
        self,
        input_ids: torch.Tensor,
        multimodal_embeddings: MultiModalEmbeddings | None = None,
    ) -> torch.Tensor:
        inputs_embeds = self.language_model.get_input_embeddings(input_ids)
        if multimodal_embeddings is not None and len(multimodal_embeddings) != 0:
            inputs_embeds = merge_multimodal_embeddings(
                inputs_embeds,
                multimodal_embeddings,
                input_ids == self.config.image_token_id,
            )
        return inputs_embeds

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
        **kwargs: object,
    ) -> torch.Tensor | IntermediateTensors:
        # V1 precomputes multimodal ``inputs_embeds`` via ``embed_multimodal`` /
        # ``get_input_embeddings``, so the model forward only threads them into
        # the language model (image + video share the same embedding merge).
        if intermediate_tensors is not None:
            inputs_embeds = None

        return self.language_model.model(
            input_ids=input_ids,
            positions=positions,
            intermediate_tensors=intermediate_tensors,
            inputs_embeds=inputs_embeds,
        )

    def compute_logits(self, hidden_states: torch.Tensor):
        return self.language_model.compute_logits(hidden_states)

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

    def get_mm_mapping(self) -> MultiModelKeys:
        return MultiModelKeys.from_string_field(
            language_model="language_model",
            connector="visual.merger.",
            tower_model="visual.",
        )

@VIDEO_LOADER_REGISTRY.register(
    "llava_onevision2",
    video_processor="LlavaOnevision2VideoProcessor",
)
class LlavaOnevision2VideoBackend(VideoBackend):
    """Frame-sampling backend for LLaVA-OneVision-2.

    Selected automatically for OV2 via the ``video_processor`` binding
    (``video_processor_type == "LlavaOnevision2VideoProcessor"`` in the model's
    ``video_preprocessor_config.json``). Decoding uses the inherited OpenCV /
    PyAV codecs; only the sampling index policy is overridden to match qwen.
    """

    _sampling_suffix = "_llava_onevision2"

    # OV2 hf-chat reference sampling constants (mirror the validated adapter).
    _FPS = _DEFAULT_FPS
    _MAX_FRAMES = _DEFAULT_MAX_FRAMES
    _MIN_FRAMES = _OV2_FPS_MIN_FRAMES

    @classmethod
    def compute_frames_index_to_sample(
        cls,
        source: VideoSourceMetadata,
        target: VideoTargetMetadata,
        **kwargs,
    ) -> list[int]:
        total = int(source.total_frames_num)
        if total <= 0:
            return []
        video_fps = float(source.original_fps)
        # Honor caller-provided sampling targets (via ``--media-io-kwargs`` →
        # ``VideoTargetMetadata``) so benchmarks can override the conservative
        # defaults (e.g. VSI-Bench needs max_frames=128). Fall back to the OV2
        # hf-chat reference constants when the target leaves a field unset
        # (sentinel ``<= 0``).
        target_fps = float(target.fps) if target.fps > 0 else cls._FPS
        target_max_frames = (
            int(target.num_frames) if target.num_frames > 0 else cls._MAX_FRAMES
        )
        n = _ov2_smart_nframes(
            total,
            video_fps,
            fps=target_fps,
            min_frames=cls._MIN_FRAMES,
            max_frames=target_max_frames,
        )
        # qwen uses linspace().round() (NOT the floor cast used by the base
        # uniform backend), so replicate the rounding exactly.
        idx = np.linspace(0, total - 1, n).round().astype(int).tolist()
        # smart_nframes floors to FRAME_FACTOR so ``n`` is even; guard anyway
        # since OV2's vision tower (temporal merge = 2) requires even counts.
        if len(idx) % _OV2_FRAME_FACTOR != 0:
            idx.append(idx[-1])
        return idx

class LlavaOnevision2VisionAttn(nn.Module):
    """Vision self-attention with windowed cu_seqlens.

    The HF checkpoint ships a *fused* qkv linear (``self_attn.qkv``), so
    we load directly into ``QKVParallelLinear`` with no stacked_params
    mapping. (Compare OV1.5, whose checkpoint had separate q/k/v.)
    """

    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        projection_size: int,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
        use_data_parallel: bool = False,
    ) -> None:
        super().__init__()
        if quant_config is not None:
            raise RuntimeError("LLaVAOneVision2 does not support quantization")

        self.tp_size = (
            1
            if use_data_parallel
            else parallel_state.get_tensor_model_parallel_world_size()
        )
        self.tp_rank = parallel_state.get_tensor_model_parallel_rank()
        self.num_heads = num_heads
        self.hidden_size_per_attn_head = dist_utils.divide(projection_size, num_heads)
        self.num_attn_heads_per_partition = dist_utils.divide(num_heads, self.tp_size)

        self.qkv = QKVParallelLinear(
            hidden_size=embed_dim,
            head_size=self.hidden_size_per_attn_head,
            total_num_heads=num_heads,
            total_num_kv_heads=num_heads,
            bias=True,
            quant_config=quant_config,
            prefix=f"{prefix}.qkv",
            disable_tp=use_data_parallel,
        )

        self.proj = RowParallelLinear(
            input_size=projection_size,
            output_size=embed_dim,
            quant_config=quant_config,
            prefix=f"{prefix}.proj",
            disable_tp=use_data_parallel,
        )

        self.attn = MMEncoderAttention(
            num_heads=self.num_attn_heads_per_partition,
            head_size=self.hidden_size_per_attn_head,
            scale=self.hidden_size_per_attn_head**-0.5,
            prefix=f"{prefix}.attn",
        )

    @staticmethod
    def _rotate_half_interleaved(x: torch.Tensor) -> torch.Tensor:
        """OV2-specific interleaved rotate_half.

        Pairs adjacent dims: (x[::2], x[1::2]) -> (-x[1::2], x[::2]).
        NOT compatible with the split-half rotate used in OV1.5/LLaMA.
        """
        x_even = x[..., 0::2]
        x_odd = x[..., 1::2]
        out = torch.stack((-x_odd, x_even), dim=-1)
        return out.flatten(-2)

    def _apply_rotary_pos_embed(
        self, t: torch.Tensor, freqs: torch.Tensor
    ) -> torch.Tensor:
        # freqs is [seq_len, half]; cat([f,f],-1) pair-repeat layout
        # matches the interleaved rotate above.
        orig_dtype = t.dtype
        t = t.float()
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = emb.cos().unsqueeze(-2).float()
        sin = emb.sin().unsqueeze(-2).float()
        t = (t * cos) + (self._rotate_half_interleaved(t) * sin)
        return t.to(orig_dtype)

    def forward(
        self,
        x: torch.Tensor,
        cu_seqlens: torch.Tensor,
        rotary_pos_emb: torch.Tensor,
        max_seqlen: torch.Tensor | None = None,
        sequence_lengths: torch.Tensor | None = None,
    ) -> torch.Tensor:
        x, _ = self.qkv(x)
        seq_len = x.shape[0]
        # QKVParallelLinear packs q/k/v along the last dim as
        # [q_heads, k_heads, v_heads]; view splits the three sections, each
        # holding this partition's heads (no all-gather: MMEncoderAttention
        # runs per-partition and ``proj`` reduces across TP ranks).
        qkv = x.view(
            seq_len,
            3,
            self.num_attn_heads_per_partition,
            self.hidden_size_per_attn_head,
        )
        q, k, v = qkv.unbind(1)

        if rotary_pos_emb is not None:
            # OV2 uses interleaved RoPE on the raw freqs, applied outside
            # MMEncoderAttention (which is rotary-agnostic).
            q = self._apply_rotary_pos_embed(q, rotary_pos_emb)
            k = self._apply_rotary_pos_embed(k, rotary_pos_emb)

        # Add a leading batch dim (b=1) for MMEncoderAttention, which expects
        # (batch, seq_len, num_heads, head_size) and windows via cu_seqlens.
        output = self.attn(
            query=q.unsqueeze(0),
            key=k.unsqueeze(0),
            value=v.unsqueeze(0),
            cu_seqlens=cu_seqlens,
            max_seqlen=max_seqlen,
            sequence_lengths=sequence_lengths,
        )
        output = output.reshape(seq_len, -1)

        output, _ = self.proj(output)
        return output

class LlavaOnevision2VisionRotaryEmbedding(nn.Module):
    """3-D rotary frequency constructor with 4:6:6 (T:H:W) split.

    Mirrors ``VisionRotaryEmbedding`` in the HF reference
    (``modeling_llava_onevision2.py`` L79-L210). The three ``inv_freq_*``
    buffers are non-persistent — they are *not* in the checkpoint and must
    be reconstructed at module init time (which we do here).

    Public entry points used by the vision tower:
      * ``forward_from_positions(patch_positions)`` — per-patch (t,h,w)
        positions → per-token freqs [N, half].
    """

    def __init__(self, head_dim: int, theta: float = 10000.0) -> None:
        super().__init__()
        assert head_dim % 2 == 0, "head_dim must be even"
        assert head_dim % 16 == 0, "head_dim must be divisible by 16 (4:6:6)"
        half = head_dim // 2
        assert half % 16 == 0, "head_dim//2 must be divisible by 16"

        self.head_dim = head_dim
        self.half = half
        self.base = float(theta)

        unit = half // 16
        self.t_size = 4 * unit
        self.h_size = 6 * unit
        self.w_size = 6 * unit
        assert self.t_size + self.h_size + self.w_size == half

        self.register_buffer(
            "inv_freq_t",
            1.0
            / (
                self.base
                ** (torch.arange(self.t_size, dtype=torch.float32) / self.t_size)
            ),
            persistent=False,
        )
        self.register_buffer(
            "inv_freq_h",
            1.0
            / (
                self.base
                ** (torch.arange(self.h_size, dtype=torch.float32) / self.h_size)
            ),
            persistent=False,
        )
        self.register_buffer(
            "inv_freq_w",
            1.0
            / (
                self.base
                ** (torch.arange(self.w_size, dtype=torch.float32) / self.w_size)
            ),
            persistent=False,
        )

    def forward_from_positions(self, patch_positions: torch.Tensor) -> torch.Tensor:
        """[N, 3] (t,h,w) int → [N, half] float frequencies."""
        device = patch_positions.device
        inv_t = self.inv_freq_t.to(device=device)
        inv_h = self.inv_freq_h.to(device=device)
        inv_w = self.inv_freq_w.to(device=device)

        t_pos = patch_positions[:, 0].float()
        h_pos = patch_positions[:, 1].float()
        w_pos = patch_positions[:, 2].float()

        ft = torch.outer(t_pos, inv_t)
        fh = torch.outer(h_pos, inv_h)
        fw = torch.outer(w_pos, inv_w)
        return torch.cat([ft, fh, fw], dim=-1)

class LlavaOnevision2VisionTower(nn.Module):
    """OV2 vision tower (no CLS token, 3-D RoPE, windowed attention).

    Module attribute names mirror HF checkpoint names verbatim so the
    WeightsMapper only needs prefix rewrites (no substring rules, which would
    otherwise collide with the Qwen3 text-path ``self_attn`` modules):
      visual.embeddings.patch_embedding
      visual.layernorm_pre
      visual.encoder.layers.{i}.self_attn.{qkv,proj}
      visual.encoder.layers.{i}.layer_norm{1,2}
      visual.encoder.layers.{i}.mlp.fc{1,2}
      visual.merger.{ln_q, mlp.{0,2}}
      visual.rotary_pos_emb (non-persistent inv_freq buffers)
    """

    def __init__(
        self,
        vision_config,
        text_hidden_size: int,
        norm_eps: float = 1e-6,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
        use_data_parallel: bool = False,
    ) -> None:
        super().__init__()
        if quant_config is not None:
            raise RuntimeError("LLaVAOneVision2 does not support quantization")

        patch_size = vision_config.patch_size
        spatial_merge_size = vision_config.spatial_merge_size
        in_channels = getattr(vision_config, "num_channels", 3)
        hidden_size = vision_config.hidden_size
        embed_dim = hidden_size
        depth = vision_config.num_hidden_layers
        num_heads = vision_config.num_attention_heads
        mlp_hidden_dim = vision_config.intermediate_size
        frame_windows_size = getattr(vision_config, "frame_windows_size", 4)
        rope_theta = getattr(vision_config, "rope_theta", 10000.0)

        self.spatial_merge_size = spatial_merge_size
        self.frame_windows_size = int(frame_windows_size)
        self.num_heads = num_heads
        self.embed_dim = embed_dim
        self.head_dim = embed_dim // num_heads
        self.use_data_parallel = use_data_parallel
        self.tp_size = (
            1
            if use_data_parallel
            else parallel_state.get_tensor_model_parallel_world_size()
        )

        self.embeddings = LlavaOnevision2VisionEmbeddings(
            patch_size=patch_size, in_channels=in_channels, embed_dim=embed_dim
        )
        self.layernorm_pre = nn.LayerNorm(embed_dim, eps=norm_eps)

        self.rotary_pos_emb = LlavaOnevision2VisionRotaryEmbedding(
            self.head_dim, theta=rope_theta
        )

        self.encoder = nn.Module()
        self.encoder.layers = nn.ModuleList(
            [
                LlavaOnevision2VisionTowerBlock(
                    dim=embed_dim,
                    num_heads=num_heads,
                    mlp_hidden_dim=mlp_hidden_dim,
                    norm_eps=norm_eps,
                    quant_config=quant_config,
                    prefix=f"{prefix}.encoder.layers.{i}",
                    use_data_parallel=use_data_parallel,
                )
                for i in range(depth)
            ]
        )

        self.merger = LlavaOnevision2PatchMerger(
            d_model=text_hidden_size,
            context_dim=embed_dim,
            spatial_merge_size=spatial_merge_size,
            norm_eps=norm_eps,
            quant_config=quant_config,
            prefix=f"{prefix}.merger",
            use_data_parallel=use_data_parallel,
        )

        # Vision attention backend; mirrors the one MMEncoderAttention picks
        # internally so the cu_seqlens / max_seqlen metadata computed below
        # matches the kernel actually used.
        self.attn_backend = get_vit_attn_backend(
            head_size=self.head_dim, dtype=torch.get_default_dtype()
        )

    @property
    def dtype(self) -> torch.dtype:
        return self.embeddings.patch_embedding.weight.dtype

    @property
    def device(self) -> torch.device:
        return self.embeddings.patch_embedding.weight.device

    def _build_window_cu_seqlens(self, grid_thw: torch.Tensor) -> np.ndarray:
        """Build cu_seqlens that chunk each sample's T-axis into windows of
        ``frame_windows_size`` frames.

        Returns an int32 ``np.ndarray`` of shape [num_windows+1] (the
        canonical prefix-sum format). Backend-specific transforms are applied
        afterwards via ``MMEncoderAttention.maybe_recompute_cu_seqlens``.
        """
        win = self.frame_windows_size
        chunk_lengths: list[int] = []
        for row in grid_thw.tolist():
            t, h, w = int(row[0]), int(row[1]), int(row[2])
            per_frame = h * w
            t_remaining = t
            while t_remaining > 0:
                this_t = min(win, t_remaining)
                chunk_lengths.append(this_t * per_frame)
                t_remaining -= this_t
        cu = np.concatenate(
            [
                np.zeros(1, dtype=np.int32),
                np.array(chunk_lengths, dtype=np.int32).cumsum(dtype=np.int32),
            ]
        )
        return cu

    def forward(
        self,
        pixel_values: torch.Tensor,
        grid_thw: torch.Tensor,
        patch_positions: torch.Tensor,
    ) -> torch.Tensor:
        x = pixel_values.to(device=self.device, dtype=self.dtype)
        x = self.embeddings(x)
        x = self.layernorm_pre(x)

        rotary_pos_emb = self.rotary_pos_emb.forward_from_positions(
            patch_positions.to(self.device)
        )

        # Build window cu_seqlens, then derive backend-specific attention
        # metadata (passthrough for FA/SDPA; transformed for FlashInfer).
        cu_seqlens_np = self._build_window_cu_seqlens(grid_thw)
        sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
            self.attn_backend, cu_seqlens_np, self.device
        )
        max_seqlen = torch.tensor(
            MMEncoderAttention.compute_max_seqlen(self.attn_backend, cu_seqlens_np),
            dtype=torch.int32,
        )
        cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
            self.attn_backend,
            cu_seqlens_np,
            self.embed_dim,
            self.tp_size,
            self.device,
        )

        for blk in self.encoder.layers:
            x = blk(
                x,
                cu_seqlens=cu_seqlens,
                rotary_pos_emb=rotary_pos_emb,
                max_seqlen=max_seqlen,
                sequence_lengths=sequence_lengths,
            )

        return self.merger(x, patch_positions=patch_positions)