Add scoped vision export support for LLaVA-OneVision-1.5

Prepare OV1.5 frame geometry before tracing and register its vision
attention class with the existing exporter adapter. Restrict the new
input preparation branch to the OV1.5 vision model type.

Test static and dynamic export with learned class tokens, isolated
frames and aspect-ratio changes at fixed patch counts. Document the
uniform-patch-count and fixed-frame-count export constraints.

Author: ZijunLi402

docs/source/en/model_doc/llava_onevision1_5.md

@@ -40,6 +40,10 @@ output = model.generate(**inputs, max_new_tokens=64)
 processor.batch_decode(output[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)
 ```
 
+## Export
+
+The vision encoder can be exported with `DynamoExporter`, which prepares grid metadata before tracing. The shared vision attention export adapter requires equal patch counts per frame. The number of frames and their patch counts are specialized to the export inputs; export a new graph when these change.
+
 ## LlavaOnevision15Config
 
 [[autodoc]] LlavaOnevision15Config

src/transformers/exporters/exporter_dynamo.py

@@ -388,6 +388,7 @@ def _to_batched(t):
     # Combined `qkv` + `(cos, sin)` rotary + `.proj`
     "transformers.models.qwen2_vl.modeling_qwen2_vl.VisionAttention.forward",
     "transformers.models.qwen2_5_vl.modeling_qwen2_5_vl.Qwen2_5_VLVisionAttention.forward",
+    "transformers.models.llava_onevision1_5.modeling_llava_onevision1_5.LlavaOnevision15VisionAttention.forward",
     "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionAttention.forward",
     "transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe.Qwen3VLMoeVisionAttention.forward",
     "transformers.models.qwen3_5.modeling_qwen3_5.Qwen3_5VisionAttention.forward",

src/transformers/exporters/utils.py

@@ -519,6 +519,9 @@ def _prepare_grid_thw_vision_inputs(model: torch.nn.Module, inputs: dict[str, An
     include_temporal = _find_submodule_attr(model, "axis_dim") is not None
     inputs["position_ids"] = get_vision_position_ids(grid_thw, spatial_merge_size, include_temporal=include_temporal)
 
+    if model.config.model_type == "llava_onevision1_5_vision":
+        inputs["frame_lengths"] = module.get_vision_frame_lengths(inputs["cu_seqlens"])
+
     window_size = _find_submodule_attr(model, "window_size")
     patch_size = _find_submodule_attr(model, "patch_size")
     if window_size is not None and patch_size is not None:

src/transformers/models/llava_onevision1_5/modeling_llava_onevision1_5.py

@@ -40,7 +40,7 @@
     merge_with_config_defaults,
 )
 from ...utils.output_capturing import capture_outputs
-from ...vision_utils import get_vision_position_ids
+from ...vision_utils import get_vision_attention_seqlens, get_vision_position_ids
 from ..auto import AutoModel
 from .configuration_llava_onevision1_5 import LlavaOnevision15Config, LlavaOnevision15VisionConfig
 
@@ -334,6 +334,13 @@ def _init_weights(self, module):
             init.normal_(module.class_pos_emb, mean=0.0, std=std)
 
 
+def get_vision_frame_lengths(cu_seqlens: torch.Tensor, kwargs: dict | None = None) -> list[int]:
+    """Get patch counts per frame, or consume geometry prepared before graph capture."""
+    if kwargs is not None and (frame_lengths := kwargs.pop("frame_lengths", None)) is not None:
+        return frame_lengths
+    return (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
+
+
 @auto_docstring
 class LlavaOnevision15VisionModel(LlavaOnevision15PreTrainedModel):
     config: LlavaOnevision15VisionConfig
@@ -368,22 +375,23 @@ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs)
             Temporal, height and width dimensions of the patch grid for each image or video.
         """
         hidden_states = self.patch_embed(hidden_states)
-        position_ids = get_vision_position_ids(grid_thw, self.spatial_merge_size)
+        position_ids = get_vision_position_ids(grid_thw, self.spatial_merge_size, kwargs=kwargs)
         cos, sin = self.rotary_pos_emb(hidden_states, position_ids)
-        lengths = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0])
-        frame_lengths = lengths.tolist()
+        cu_seqlens, max_seqlen = get_vision_attention_seqlens(grid_thw, self.config, kwargs=kwargs)
+        frame_lengths = get_vision_frame_lengths(cu_seqlens, kwargs=kwargs)
+        cu_seqlens = cu_seqlens + torch.arange(cu_seqlens.shape[0], device=cu_seqlens.device, dtype=cu_seqlens.dtype)
+        if max_seqlen is not None:
+            max_seqlen += 1
         cls_hidden = self.class_embedding.to(hidden_states.dtype).unsqueeze(0)
         cls_freqs = torch.cat((self.class_pos_emb, self.class_pos_emb), dim=-1).to(cos.dtype)
         cls_cos, cls_sin = cls_freqs.cos().to(cos.dtype), cls_freqs.sin().to(sin.dtype)
         hidden_states, cos, sin = [
             torch.cat([torch.cat((cls, frame), dim=0) for frame in torch.split(value, frame_lengths)], dim=0)
             for value, cls in ((hidden_states, cls_hidden), (cos, cls_cos), (sin, cls_sin))
         ]
-        cu_seqlens = nn.functional.pad((lengths + 1).cumsum(0, dtype=torch.int32), (1, 0))
-
         hidden_states = self.pre_layernorm(hidden_states)
         for block in self.blocks:
-            hidden_states = block(hidden_states, cu_seqlens, cos, sin, **kwargs)
+            hidden_states = block(hidden_states, cu_seqlens, cos, sin, max_seqlen=max_seqlen, **kwargs)
 
         hidden_states = torch.cat(
             [frame[1:] for frame in torch.split(hidden_states, [length + 1 for length in frame_lengths])], dim=0

src/transformers/models/llava_onevision1_5/modular_llava_onevision1_5.py

@@ -27,7 +27,7 @@
 from ...utils import auto_docstring
 from ...utils.generic import merge_with_config_defaults
 from ...utils.output_capturing import capture_outputs
-from ...vision_utils import get_vision_position_ids
+from ...vision_utils import get_vision_attention_seqlens, get_vision_position_ids
 from ..auto import AutoConfig, AutoModel
 from ..qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig, Qwen2_5_VLVisionConfig
 from ..qwen2_5_vl.modeling_qwen2_5_vl import (
@@ -114,6 +114,13 @@ def __post_init__(self, **kwargs):
             raise ValueError("vision_config.out_hidden_size must match text_config.hidden_size")
 
 
+def get_vision_frame_lengths(cu_seqlens: torch.Tensor, kwargs: dict | None = None) -> list[int]:
+    """Get patch counts per frame, or consume geometry prepared before graph capture."""
+    if kwargs is not None and (frame_lengths := kwargs.pop("frame_lengths", None)) is not None:
+        return frame_lengths
+    return (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
+
+
 class LlavaOnevision15VisionPatchEmbed(nn.Module):
     def __init__(self, config: LlavaOnevision15VisionConfig):
         super().__init__()
@@ -239,22 +246,23 @@ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs)
             Temporal, height and width dimensions of the patch grid for each image or video.
         """
         hidden_states = self.patch_embed(hidden_states)
-        position_ids = get_vision_position_ids(grid_thw, self.spatial_merge_size)
+        position_ids = get_vision_position_ids(grid_thw, self.spatial_merge_size, kwargs=kwargs)
         cos, sin = self.rotary_pos_emb(hidden_states, position_ids)
-        lengths = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0])
-        frame_lengths = lengths.tolist()
+        cu_seqlens, max_seqlen = get_vision_attention_seqlens(grid_thw, self.config, kwargs=kwargs)
+        frame_lengths = get_vision_frame_lengths(cu_seqlens, kwargs=kwargs)
+        cu_seqlens = cu_seqlens + torch.arange(cu_seqlens.shape[0], device=cu_seqlens.device, dtype=cu_seqlens.dtype)
+        if max_seqlen is not None:
+            max_seqlen += 1
         cls_hidden = self.class_embedding.to(hidden_states.dtype).unsqueeze(0)
         cls_freqs = torch.cat((self.class_pos_emb, self.class_pos_emb), dim=-1).to(cos.dtype)
         cls_cos, cls_sin = cls_freqs.cos().to(cos.dtype), cls_freqs.sin().to(sin.dtype)
         hidden_states, cos, sin = [
             torch.cat([torch.cat((cls, frame), dim=0) for frame in torch.split(value, frame_lengths)], dim=0)
             for value, cls in ((hidden_states, cls_hidden), (cos, cls_cos), (sin, cls_sin))
         ]
-        cu_seqlens = nn.functional.pad((lengths + 1).cumsum(0, dtype=torch.int32), (1, 0))
-
         hidden_states = self.pre_layernorm(hidden_states)
         for block in self.blocks:
-            hidden_states = block(hidden_states, cu_seqlens, cos, sin, **kwargs)
+            hidden_states = block(hidden_states, cu_seqlens, cos, sin, max_seqlen=max_seqlen, **kwargs)
 
         hidden_states = torch.cat(
             [frame[1:] for frame in torch.split(hidden_states, [length + 1 for length in frame_lengths])], dim=0

tests/models/llava_onevision1_5/test_modeling_llava_onevision1_5.py

@@ -36,6 +36,7 @@
     require_flash_attn,
     require_torch,
     require_torch_gpu,
+    require_torch_greater_or_equal,
     slow,
     torch_device,
 )
@@ -338,6 +339,42 @@ def test_mismatching_num_image_tokens(self):
 class LlavaOnevision15VisionModelTest(unittest.TestCase):
     all_model_classes = (LlavaOnevision15VisionModel,) if is_torch_available() else ()
 
+    @slow
+    @require_torch_greater_or_equal("2.11")
+    def test_export_preserves_class_tokens_and_frame_isolation(self):
+        from transformers.exporters import DynamoConfig, DynamoExporter
+        from transformers.exporters.utils import prepare_for_export
+
+        config = LlavaOnevision15VisionConfig(
+            hidden_size=32, out_hidden_size=32, intermediate_size=48, depth=2, num_heads=4
+        )
+        model = LlavaOnevision15VisionModel(config).to(torch_device).eval()
+        # One image and two video frames, with 16 patches per frame.
+        inputs = {
+            "hidden_states": floats_tensor([48, 3 * 14 * 14]),
+            "grid_thw": torch.tensor([[1, 4, 4], [2, 4, 4]], device=torch_device),
+        }
+        for dynamic in (False, True):
+            with self.subTest(dynamic=dynamic):
+                sample = copy.deepcopy(inputs)
+                exported = DynamoExporter().export(model, sample, config=DynamoConfig(dynamic=dynamic))
+                # Replay with new pixels and aspect ratio, retaining 16 patches per frame.
+                sample["hidden_states"] = floats_tensor([48, 3 * 14 * 14])
+                sample["grid_thw"] = torch.tensor([[1, 2, 8], [2, 2, 8]], device=torch_device)
+                model, sample, _ = prepare_for_export(model, sample)
+                with torch.no_grad():
+                    actual = exported.module()(**sample)
+                    expected = model(sample["hidden_states"], sample["grid_thw"])
+                    separate = torch.cat(
+                        [
+                            model(frame, sample["grid_thw"][:1]).pooler_output
+                            for frame in sample["hidden_states"].split(16)
+                        ]
+                    )
+                torch.testing.assert_close(actual.last_hidden_state, expected.last_hidden_state)
+                torch.testing.assert_close(actual.pooler_output, expected.pooler_output)
+                torch.testing.assert_close(actual.pooler_output, separate)
+
     def test_standalone_save_load_and_frame_isolation(self):
         config = LlavaOnevision15VisionConfig(
             hidden_size=32, out_hidden_size=32, intermediate_size=48, depth=2, num_heads=4