optimum-rbln 0.9.3rc0__py3-none-any.whl → 0.9.5a4__py3-none-any.whl

optimum/rbln/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py

@@ -17,7 +17,13 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple, Union
 
 import torch
-from transformers import AutoModelForVision2Seq, PretrainedConfig, PreTrainedModel, Qwen2_5_VLForConditionalGeneration
+from transformers import (
+    AutoModelForVision2Seq,
+    PretrainedConfig,
+    PreTrainedModel,
+    Qwen2_5_VLConfig,
+    Qwen2_5_VLForConditionalGeneration,
+)
 from transformers.modeling_utils import no_init_weights
 from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
     Qwen2_5_VisionPatchEmbed,
@@ -30,8 +36,8 @@ from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
 from ....configuration_utils import RBLNCompileConfig
 from ....modeling import RBLNModel
 from ....utils.logging import get_logger
-from ...modeling_outputs import RBLNDecoderOnlyOutput
-from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyModelForCausalLM
+from ...modeling_outputs import RBLNDecoderOnlyOutput, _validate_output_hidden_states
+from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyModel, RBLNDecoderOnlyModelForCausalLM
 from .configuration_qwen2_5_vl import (
     RBLNQwen2_5_VisionTransformerPretrainedModelConfig,
     RBLNQwen2_5_VLForConditionalGenerationConfig,
@@ -42,7 +48,7 @@ from .qwen2_5_vl_architecture import Qwen2_5_VisionTransformerWrapper, Qwen2_5_V
 logger = get_logger(__name__)
 
 if TYPE_CHECKING:
-    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer
 
 
 class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
@@ -55,6 +61,7 @@ class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
     """
 
     auto_model_class = None
+    _supports_non_fp32 = True
 
     def __post_init__(self, **kwargs):
         self.transformer = self.model[0]
@@ -88,10 +95,10 @@ class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
         torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
 
     @classmethod
-    def wrap_model_if_needed(
+    def _wrap_model_if_needed(
         cls, model: "PreTrainedModel", rbln_config: RBLNQwen2_5_VisionTransformerPretrainedModelConfig
     ):
-        return Qwen2_5_VisionTransformerWrapper(model).eval()
+        return Qwen2_5_VisionTransformerWrapper(model, rbln_config).eval()
 
     def __getattr__(self, __name: str) -> Any:
         def redirect(func):
@@ -111,10 +118,10 @@ class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
         model_config: "PretrainedConfig" = None,
         rbln_config: Optional[RBLNQwen2_5_VisionTransformerPretrainedModelConfig] = None,
     ) -> RBLNQwen2_5_VisionTransformerPretrainedModelConfig:
-        window_size = getattr(model_config, "window_size")
-        patch_size = getattr(model_config, "patch_size")
-        hidden_size = getattr(model_config, "hidden_size")
-        num_heads = getattr(model_config, "num_heads")
+        window_size = model_config.window_size
+        patch_size = model_config.patch_size
+        hidden_size = model_config.hidden_size
+        num_heads = model_config.num_heads
         head_dim = hidden_size // num_heads
         window_seq_len = (window_size // patch_size) ** 2
 
@@ -126,22 +133,22 @@ class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
                 )
 
             input_info = [
-                ("hidden_states", [max_seq_len, hidden_size], "float32"),
-                ("full_attn_masks", [1, 1, max_seq_len, max_seq_len], "float32"),
+                ("hidden_states", [max_seq_len, hidden_size], rbln_config.dtype),
+                ("full_attn_masks", [1, 1, max_seq_len, max_seq_len], rbln_config.dtype),
                 (
                     "window_attn_masks",
                     [max_seq_len // window_seq_len, 1, window_seq_len, window_seq_len],
-                    "float32",
+                    rbln_config.dtype,
                 ),
                 (
                     "cos",
                     [1, 1, max_seq_len, head_dim],
-                    "float32",
+                    rbln_config.dtype,
                 ),
                 (
                     "sin",
                     [1, 1, max_seq_len, head_dim],
-                    "float32",
+                    rbln_config.dtype,
                 ),
             ]
             input_infos.append(input_info)
@@ -203,7 +210,7 @@ class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
             1,
             window_seq_len,
             window_seq_len,
-            dtype=torch.float32,
+            dtype=hidden_states.dtype,
         )
         for i, valid_len in enumerate(window_valid_lengths):
             if valid_len < window_seq_len:
@@ -242,7 +249,7 @@ class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
             1,
             max_seq_len,
             max_seq_len,
-            dtype=torch.float32,
+            dtype=hidden_state_padded.dtype,
         )
         for i, valid_len in enumerate(window_valid_lengths):
             start = i * window_seq_len
@@ -253,7 +260,7 @@ class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
         return hidden_state_full_padded, cos_full_padded, sin_full_padded, full_attn_masks
 
     def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor:
-        hidden_states = self.patch_embed(hidden_states)
+        hidden_states = self.patch_embed(hidden_states).to(self.rbln_config.dtype)
         rotary_pos_emb = self.rot_pos_emb(grid_thw)
         window_index, cu_window_seqlens = self.get_window_index(grid_thw)
         cu_window_seqlens = torch.tensor(
@@ -270,7 +277,7 @@ class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
         rotary_pos_emb = rotary_pos_emb[window_index, :, :]
         rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
         emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
-        position_embeddings = (emb.cos(), emb.sin())
+        position_embeddings = (emb.cos().to(self.rbln_config.dtype), emb.sin().to(self.rbln_config.dtype))
 
         cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
             dim=0,
@@ -294,10 +301,10 @@ class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
             try:
                 ws_index = torch.searchsorted(self.max_seq_lens, window_padded_len).item()
                 max_seq_len = self.max_seq_lens[ws_index]
-            except Exception:
+            except Exception as e:
                 raise ValueError(
                     f"Required seq_len({window_padded_len}) is larger than available max_seq_lens({self.max_seq_lens.tolist()})."
-                )
+                ) from e
 
             # Padding for Window Attention Layers
             hidden_state_padded, cos_padded, sin_padded, window_attn_masks, window_valid_lengths = (
@@ -338,67 +345,47 @@ class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
         return hidden_states
 
 
-class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
-    """
-    RBLNQwen2_5_VLForConditionalGeneration is a multi-modal model that integrates vision and language processing capabilities,
-    optimized for RBLN NPUs. It is designed for conditional generation tasks that involve both image and text inputs.
-
-    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. Check the superclass documentation for the generic methods the library implements for all its models.
-
-    Important Note:
-        This model includes a Large Language Model (LLM). For optimal performance, it is highly recommended to use
-        tensor parallelism for the language model. This can be achieved by using the `rbln_config` parameter in the
-        `from_pretrained` method. Refer to the `from_pretrained` documentation and the RBLNQwen2_5_VLForConditionalGenerationConfig class for details.
-
-    Examples:
-        ```python
-        from optimum.rbln import RBLNQwen2_5_VLForConditionalGeneration
-
-        model = RBLNQwen2_5_VLForConditionalGeneration.from_pretrained(
-            "Qwen/Qwen2.5-VL-7B-Instruct",
-            export=True,
-            rbln_config={
-                "visual": {
-                    "max_seq_lens": 6400,
-                    "device": 0,
-                },
-                "tensor_parallel_size": 8,
-                "kvcache_partition_len": 16_384,
-                "max_seq_len": 114_688,
-                "device": [0, 1, 2, 3, 4, 5, 6, 7],
-            },
-        )
-
-        model.save_pretrained("compiled-qwen2.5-vl-7b-instruct")
-        ```
-    """
-
-    _supports_non_fp32 = False
-
+class RBLNQwen2_5_VLModel(RBLNDecoderOnlyModel):
     auto_model_class = AutoModelForVision2Seq
+    _decoder_wrapper_cls = Qwen2_5_VL_LanguageModelWrapper
+    _use_rotary_emb = False
     _rbln_submodules = [
         {"name": "visual"},
     ]
-    _decoder_wrapper_cls = Qwen2_5_VL_LanguageModelWrapper
-    _use_rotary_emb = False
+    _config_class = Qwen2_5_VLConfig
+    _rotary_emb_class = Qwen2_5_VLRotaryEmbedding
+    _get_rope_index_func = Qwen2_5_VLModel.get_rope_index
 
     def __post_init__(self, **kwargs):
+        if hasattr(self.config, "embedding_dim"):
+            self.embedding_dim = self.config.embedding_dim
+
+        if not isinstance(self.config.text_config, PretrainedConfig):
+            self.config = self._config_class(
+                text_config=self.config.text_config, vision_config=self.config.vision_config
+            )
+
         super().__post_init__(**kwargs)
         self.visual = self.rbln_submodules[0]
-        self.mrope_section = self.config.rope_scaling["mrope_section"]
-        self.rotary_emb = Qwen2_5_VLRotaryEmbedding(self.config)
-        self.rope_deltas = torch.zeros(self.rbln_config.batch_size)
-
-    def can_generate(self):
-        return True
+        self.rotary_emb = self._rotary_emb_class(self.config)
+        if not self.can_generate():
+            self.block_tables = torch.arange(self.rbln_config.kvcache_num_blocks, dtype=torch.int16)
+
+    @property
+    def logits_last_dim(self):
+        if self.can_generate():
+            return self.config.vocab_size
+        else:
+            return self.embedding_dim if hasattr(self, "embedding_dim") else self.config.hidden_size
 
-    @classmethod
-    def get_pytorch_model(cls, *args, **kwargs):
-        model = super().get_pytorch_model(*args, **kwargs)
-        model.model.lm_head = model.lm_head
-        model.lm_head = None
-        del model.lm_head
-        return model
+    def _create_embedding_layer(self):
+        with no_init_weights():
+            embed_tokens = torch.nn.Embedding(
+                self.config.text_config.vocab_size,
+                self.config.text_config.hidden_size,
+                self.config.text_config.pad_token_id,
+            )
+        return embed_tokens
 
     @classmethod
     def get_input_info(
@@ -415,61 +402,25 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
             (
                 "position_emb",
                 [2, batch_size, 1, query_length, model_config.hidden_size // model_config.num_attention_heads],
-                "float32",
+                rbln_config.dtype,
             ),
         )
 
         return input_info
 
-    def prepare_inputs_for_generation(
-        self,
-        input_ids: torch.LongTensor,
-        generate_idx: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.LongTensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        pixel_values=None,
-        pixel_values_videos=None,
-        image_grid_thw=None,
-        video_grid_thw=None,
-        second_per_grid_ts=None,
-        **kwargs,
-    ):
-        model_inputs = {}
-        is_prefill_phase = generate_idx is None
-
-        if is_prefill_phase:
-            generate_idx = attention_mask.sum(dim=-1, keepdim=True).int()
-            cache_position = None
-            model_inputs.update({"input_ids": input_ids})
-        else:
-            if inputs_embeds is not None:
-                raise NotImplementedError("Specifying inputs_embeds in decoder phase is not supported.")
-
-            input_ids = input_ids[:, -1:]
-            cache_position = generate_idx
-            generate_idx = generate_idx + 1
-            model_inputs.update({"input_ids": input_ids})
-
-        model_inputs.update(
-            {
-                "attention_mask": attention_mask,
-                "cache_position": cache_position,
-                "generate_idx": generate_idx,
-                "pixel_values": pixel_values,
-                "pixel_values_videos": pixel_values_videos,
-                "image_grid_thw": image_grid_thw,
-                "video_grid_thw": video_grid_thw,
-                "second_per_grid_ts": second_per_grid_ts,
-            }
-        )
-
-        return model_inputs
-
     def _get_position_embeddings(self, hidden_states, position_ids):
         cos, sin = self.rotary_emb(hidden_states, position_ids)
-        mrope_section = self.mrope_section * 2
-        cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(1)
-        sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(1)
+        mrope_section = self.config.rope_scaling["mrope_section"] * 2
+        cos = (
+            torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1)
+            .unsqueeze(1)
+            .to(self.rbln_config.dtype)
+        )
+        sin = (
+            torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1)
+            .unsqueeze(1)
+            .to(self.rbln_config.dtype)
+        )
         return torch.stack([cos, sin])
 
     def _preprocess_prefill(
@@ -483,7 +434,7 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
         second_per_grid_ts: torch.Tensor = None,
     ):
         batch_size = input_ids.shape[0]
-        inputs_embeds = self.embed_tokens(input_ids)
+        inputs_embeds = self.embed_tokens(input_ids).to(self.rbln_config.dtype)
 
         if pixel_values is not None:
             image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
@@ -518,7 +469,7 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
         max_inputs_len = input_ids.shape[1]
 
         head_dim = getattr(self.config, "head_dim", None) or self.config.hidden_size // self.config.num_attention_heads
-        all_position_embeds = torch.zeros(2, batch_size, 1, max_inputs_len, head_dim)
+        all_position_embeds = torch.zeros(2, batch_size, 1, max_inputs_len, head_dim, dtype=self.rbln_config.dtype)
         all_rope_deltas = []
 
         image_token_id = self.config.image_token_id
@@ -532,8 +483,7 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
             vision_tokens = input_id[0][vision_start_indices + 1]
             image_nums = (vision_tokens == image_token_id).sum()
             video_nums = (vision_tokens == video_token_id).sum()
-            position_ids, rope_deltas = Qwen2_5_VLModel.get_rope_index(
-                self,
+            position_ids, rope_deltas = self._get_rope_index_func(
                 input_id,
                 image_grid_thw[image_idx : image_idx + image_nums] if image_grid_thw is not None else None,
                 video_grid_thw[video_idx : video_idx + video_nums] if video_grid_thw is not None else None,
@@ -551,6 +501,180 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
 
         return inputs_embeds, all_position_embeds, rope_deltas
 
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        pixel_values: Optional[torch.Tensor] = None,
+        pixel_values_videos: Optional[torch.FloatTensor] = None,
+        image_grid_thw: Optional[torch.LongTensor] = None,
+        video_grid_thw: Optional[torch.LongTensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        second_per_grid_ts: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> RBLNDecoderOnlyOutput:
+        inputs_embeds, position_embed, rope_deltas = self._preprocess_prefill(
+            input_ids,
+            attention_mask,
+            pixel_values,
+            pixel_values_videos,
+            image_grid_thw,
+            video_grid_thw,
+            second_per_grid_ts,
+        )
+
+        self.rope_deltas = rope_deltas
+        batch_size, seq_len = inputs_embeds.shape[:2]
+
+        output_hidden_states = _validate_output_hidden_states(output_hidden_states, self.rbln_config)
+
+        all_hidden_states = (
+            tuple(
+                torch.zeros(
+                    batch_size,
+                    seq_len,
+                    self.config.hidden_size,
+                    dtype=self.rbln_config.dtype,
+                )
+                for _ in range(self.config.num_hidden_layers + 1)
+            )
+            if output_hidden_states
+            else None
+        )
+
+        logits = []
+        for b_idx in range(batch_size):
+            query_length = attention_mask[b_idx].sum(dim=-1).int().item()
+            cache_position = torch.arange(query_length, dtype=torch.int32).unsqueeze(0)
+
+            output = self.prefill_decoder(
+                inputs_embeds=inputs_embeds[b_idx : b_idx + 1],
+                attention_mask=attention_mask[b_idx] if attention_mask is not None else None,
+                cache_position=cache_position,
+                batch_idx=b_idx,
+                position_embed=position_embed[:, b_idx : b_idx + 1],
+                block_tables=self.block_tables,
+            )
+            logits.append(output.logits)
+            if self.rbln_config.output_hidden_states:
+                for l_idx in range(self.config.num_hidden_layers + 1):
+                    all_hidden_states[l_idx][b_idx].copy_(output.hidden_states[l_idx][0])
+        logits = torch.cat(logits, dim=0)
+
+        if not return_dict:
+            return_value = logits if not output_hidden_states else (logits, all_hidden_states)
+            return return_value
+        else:
+            return (
+                RBLNDecoderOnlyOutput(logits=logits, hidden_states=all_hidden_states)
+                if output_hidden_states
+                else RBLNDecoderOnlyOutput(logits=logits)
+            )
+
+
+# MRO: RBLNQwen2_5_VLForConditionalGeneration -> RBLNQwen2_5_VLModel -> RBLNDecoderOnlyModelForCausalLM -> RBLNDecoderOnlyModel -> RBLNModel
+class RBLNQwen2_5_VLForConditionalGeneration(RBLNQwen2_5_VLModel, RBLNDecoderOnlyModelForCausalLM):
+    """
+    RBLNQwen2_5_VLForConditionalGeneration is a multi-modal model that integrates vision and language processing capabilities,
+    optimized for RBLN NPUs. It is designed for conditional generation tasks that involve both image and text inputs.
+
+    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    Important Note:
+        This model includes a Large Language Model (LLM). For optimal performance, it is highly recommended to use
+        tensor parallelism for the language model. This can be achieved by using the `rbln_config` parameter in the
+        `from_pretrained` method. Refer to the `from_pretrained` documentation and the RBLNQwen2_5_VLForConditionalGenerationConfig class for details.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNQwen2_5_VLForConditionalGeneration
+
+        model = RBLNQwen2_5_VLForConditionalGeneration.from_pretrained(
+            "Qwen/Qwen2.5-VL-7B-Instruct",
+            export=True,
+            rbln_config={
+                "visual": {
+                    "max_seq_lens": 6400,
+                    "device": 0,
+                },
+                "tensor_parallel_size": 8,
+                "kvcache_partition_len": 16_384,
+                "max_seq_len": 114_688,
+                "device": [0, 1, 2, 3, 4, 5, 6, 7],
+            },
+        )
+
+        model.save_pretrained("compiled-qwen2.5-vl-7b-instruct")
+        ```
+    """
+
+    auto_model_class = AutoModelForVision2Seq
+    _decoder_wrapper_cls = Qwen2_5_VL_LanguageModelWrapper
+    _supports_non_fp32 = True
+    _use_rotary_emb = False
+    _rbln_submodules = [
+        {"name": "visual"},
+    ]
+
+    def __post_init__(self, **kwargs):
+        super().__post_init__(**kwargs)
+        self.rope_deltas = torch.zeros(self.rbln_config.batch_size)
+
+    def can_generate(self):
+        return True
+
+    @classmethod
+    def _reconstruct_model_if_needed(cls, model: "PreTrainedModel"):
+        model.model.lm_head = model.lm_head
+        return model
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        generate_idx: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        pixel_values=None,
+        pixel_values_videos=None,
+        image_grid_thw=None,
+        video_grid_thw=None,
+        second_per_grid_ts=None,
+        **kwargs,
+    ):
+        model_inputs = {}
+        is_prefill_phase = generate_idx is None
+
+        if is_prefill_phase:
+            generate_idx = attention_mask.sum(dim=-1, keepdim=True).int()
+            cache_position = None
+            model_inputs.update({"input_ids": input_ids})
+        else:
+            if inputs_embeds is not None:
+                raise NotImplementedError("Specifying inputs_embeds in decoder phase is not supported.")
+
+            input_ids = input_ids[:, -1:]
+            cache_position = generate_idx
+            generate_idx = generate_idx + 1
+            model_inputs.update({"input_ids": input_ids})
+
+        model_inputs.update(
+            {
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+                "generate_idx": generate_idx,
+                "pixel_values": pixel_values,
+                "pixel_values_videos": pixel_values_videos,
+                "image_grid_thw": image_grid_thw,
+                "video_grid_thw": video_grid_thw,
+                "second_per_grid_ts": second_per_grid_ts,
+            }
+        )
+
+        return model_inputs
+
     def _preprocess_decoder(
         self,
         input_ids: torch.LongTensor = None,
@@ -561,14 +685,14 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
                 f"Cache position size mismatch: got {cache_position.shape[0]}, expected {self.rbln_config.batch_size}."
             )
 
-        inputs_embeds = self.embed_tokens(input_ids)
+        inputs_embeds = self.embed_tokens(input_ids).to(self.rbln_config.dtype)
         position_embeds = []
         for b_idx in range(self.rbln_config.batch_size):
             delta = cache_position[b_idx] + self.rope_deltas[b_idx]
             position_ids = torch.arange(1).view(1, -1)
             position_ids = position_ids.add(delta)
             position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
-            position_embed = self._get_position_embeddings(torch.zeros(1, dtype=torch.float32), position_ids)
+            position_embed = self._get_position_embeddings(torch.zeros(1, dtype=self.rbln_config.dtype), position_ids)
             position_embeds.append(position_embed)
 
         position_embeds = torch.cat(position_embeds, dim=1)
@@ -588,8 +712,10 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
         second_per_grid_ts: Optional[torch.Tensor] = None,
         generate_idx: Optional[torch.Tensor] = None,
         return_dict: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
         **kwargs,
     ) -> RBLNDecoderOnlyOutput:
+        output_hidden_states = _validate_output_hidden_states(output_hidden_states, self.rbln_config)
         # Prefill
         if cache_position is None:
             inputs_embeds, position_embed, rope_deltas = self._preprocess_prefill(
@@ -602,8 +728,21 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
                 second_per_grid_ts,
             )
 
+            batch_size, seq_len = inputs_embeds.shape[:2]
+            all_hidden_states = (
+                tuple(
+                    torch.zeros(
+                        batch_size,
+                        seq_len,
+                        self.config.hidden_size,
+                        dtype=self.rbln_config.dtype,
+                    )
+                    for _ in range(self.config.num_hidden_layers + 1)
+                )
+                if output_hidden_states
+                else None
+            )
             self.rope_deltas = rope_deltas
-            batch_size = inputs_embeds.shape[0]
 
             logits = []
             for b_idx in range(batch_size):
@@ -617,8 +756,11 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
                     position_embed=position_embed[:, b_idx : b_idx + 1],
                 )
                 logits.append(output.logits)
+                if self.rbln_config.output_hidden_states:
+                    for l_idx in range(self.config.num_hidden_layers + 1):
+                        all_hidden_states[l_idx][b_idx].copy_(output.hidden_states[l_idx][0])
             logits = torch.cat(logits, dim=0)
-            # Decoder
+        # Decoder
         else:
             inputs_embeds, position_embed = self._preprocess_decoder(input_ids, cache_position)
             output = self.decoder(
@@ -627,11 +769,17 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
                 position_embed=position_embed,
             )
             logits = output.logits
+            all_hidden_states = output.hidden_states
 
         if not return_dict:
-            return logits, generate_idx
+            return_value = (
+                logits,
+                generate_idx if not output_hidden_states else (logits, generate_idx, all_hidden_states),
+            )
+            return return_value
         else:
             return RBLNDecoderOnlyOutput(
                 logits=logits,
                 generate_idx=generate_idx,
+                hidden_states=all_hidden_states,
             )