update

Gradio Demo Example, Incremental Prefilling and VLMEvalKit Support

deepseek_vl2/models/modeling_deepseek.py

@@ -17,7 +17,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-""" PyTorch DeepSeek model."""
+""" PyTorch DeepSeek model and compatible with both DeepSeekV2 and DeepSeekV3"""
 import math
 import warnings
 from typing import List, Optional, Tuple, Union
@@ -27,16 +27,13 @@ import torch
 import torch.nn.functional as F
 import torch.utils.checkpoint
 import torch.distributed as dist
+from einops import repeat
 from torch import nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache
-from transformers.modeling_attn_mask_utils import (
-    AttentionMaskConverter,
-    _prepare_4d_attention_mask,
-    _prepare_4d_causal_attention_mask,
-)
+from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
 from transformers.models.llama.modeling_llama import (
     LlamaAttention,
     LlamaFlashAttention2
@@ -63,12 +60,10 @@ from transformers.utils.import_utils import is_torch_fx_available
 
 from .configuration_deepseek import DeepseekV2Config
 
-
 if is_flash_attn_2_available():
     from flash_attn import flash_attn_func, flash_attn_varlen_func
     from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
 
-
 # This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
 # It means that the function will not be traced through and simply appear as a node in the graph.
 if is_torch_fx_available():
@@ -77,7 +72,6 @@ if is_torch_fx_available():
 
     _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)
 
-
 logger = logging.get_logger(__name__)
 
 _CONFIG_FOR_DOC = "DeepseekV2Config"
@@ -869,17 +863,10 @@ class DeepseekV2Attention(nn.Module):
         compressed_kv, k_pe = torch.split(
             compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
         )
+        compressed_kv = self.kv_a_layernorm(compressed_kv)
         k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2)
-        kv = (
-            self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
-            .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
-            .transpose(1, 2)
-        )
 
-        k_nope, value_states = torch.split(
-            kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1
-        )
-        kv_seq_len = value_states.shape[-2]
+        kv_seq_len = k_pe.shape[-2]
         if past_key_value is not None:
             if self.layer_idx is None:
                 raise ValueError(
@@ -888,27 +875,23 @@ class DeepseekV2Attention(nn.Module):
                     "with a layer index."
                 )
             kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
-        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
 
+        cos, sin = self.rotary_emb(q_pe, seq_len=kv_seq_len)
         q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids)
 
-        query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)
-        query_states[:, :, :, : self.qk_nope_head_dim] = q_nope
-        query_states[:, :, :, self.qk_nope_head_dim :] = q_pe
-
-        key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)
-        key_states[:, :, :, : self.qk_nope_head_dim] = k_nope
-        key_states[:, :, :, self.qk_nope_head_dim :] = k_pe
         if past_key_value is not None:
             cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
-            key_states, value_states = past_key_value.update(
-                key_states, value_states, self.layer_idx, cache_kwargs
-            )
+            compressed_kv = compressed_kv.unsqueeze(1)
+            k_pe, compressed_kv = past_key_value.update(k_pe, compressed_kv, self.layer_idx, cache_kwargs)
+            compressed_kv = compressed_kv.squeeze(1)
 
-        attn_weights = (
-            torch.matmul(query_states, key_states.transpose(2, 3)) * self.softmax_scale
-        )
+        kv_b_proj = self.kv_b_proj.weight.view(self.num_heads, -1, self.kv_lora_rank)
+        q_absorb = kv_b_proj[:, :self.qk_nope_head_dim, :]
+        out_absorb = kv_b_proj[:, self.qk_nope_head_dim:, :]
 
+        q_nope = torch.matmul(q_nope, q_absorb)
+        attn_weights = (torch.matmul(q_pe, k_pe.mT) +
+                        torch.matmul(q_nope, compressed_kv.unsqueeze(-3).mT)) * self.softmax_scale
         if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
             raise ValueError(
                 f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
@@ -925,11 +908,13 @@ class DeepseekV2Attention(nn.Module):
         # upcast attention to fp32
         attn_weights = nn.functional.softmax(
             attn_weights, dim=-1, dtype=torch.float32
-        ).to(query_states.dtype)
+        ).to(q_pe.dtype)
         attn_weights = nn.functional.dropout(
             attn_weights, p=self.attention_dropout, training=self.training
         )
-        attn_output = torch.matmul(attn_weights, value_states)
+        attn_output = torch.einsum('bhql,blc->bhqc', attn_weights, compressed_kv)
+
+        attn_output = torch.matmul(attn_output, out_absorb.mT)
 
         if attn_output.size() != (bsz, self.num_heads, q_len, self.v_head_dim):
             raise ValueError(
@@ -1034,6 +1019,7 @@ class DeepseekV2FlashAttention2(DeepseekV2Attention):
         if self.q_head_dim != self.v_head_dim:
             value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim])
 
+        # TODO: support compressed_kv for kv_cache (instead of key_states, value_states) in flash_attention version
         if past_key_value is not None:
             cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
             key_states, value_states = past_key_value.update(
@@ -1494,6 +1480,7 @@ class DeepseekV2Model(DeepseekV2PreTrainedModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None
     ) -> Union[Tuple, BaseModelOutputWithPast]:
         output_attentions = (
             output_attentions
@@ -1668,17 +1655,18 @@ class DeepseekV2ForCausalLM(DeepseekV2PreTrainedModel):
         output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
     )
     def forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[List[torch.FloatTensor]] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
+            self,
+            input_ids: torch.LongTensor = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[List[torch.FloatTensor]] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+            cache_position: Optional[torch.LongTensor] = None
     ) -> Union[Tuple, CausalLMOutputWithPast]:
         r"""
         Args:
@@ -1730,6 +1718,7 @@ class DeepseekV2ForCausalLM(DeepseekV2PreTrainedModel):
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
+            cache_position=cache_position
         )
 
         hidden_states = outputs[0]
@@ -1762,13 +1751,14 @@ class DeepseekV2ForCausalLM(DeepseekV2PreTrainedModel):
         )
 
     def prepare_inputs_for_generation(
-        self,
-        input_ids,
-        past_key_values=None,
-        attention_mask=None,
-        inputs_embeds=None,
-        **kwargs,
+            self,
+            input_ids,
+            past_key_values=None,
+            attention_mask=None,
+            inputs_embeds=None,
+            **kwargs,
     ):
+        past_length = 0
         if past_key_values is not None:
             if isinstance(past_key_values, Cache):
                 cache_length = past_key_values.get_seq_length()
@@ -1780,13 +1770,10 @@ class DeepseekV2ForCausalLM(DeepseekV2PreTrainedModel):
 
             # Keep only the unprocessed tokens:
             # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
-            # some of the inputs are exclusivelly passed as part of the cache (e.g. when passing input_embeds as
+            # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
             # input)
-            if (
-                attention_mask is not None
-                and attention_mask.shape[1] > input_ids.shape[1]
-            ):
-                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
+            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
+                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length):]
             # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
             # input_ids based on the past_length.
             elif past_length < input_ids.shape[1]:
@@ -1795,9 +1782,9 @@ class DeepseekV2ForCausalLM(DeepseekV2PreTrainedModel):
 
             # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
             if (
-                max_cache_length is not None
-                and attention_mask is not None
-                and cache_length + input_ids.shape[1] > max_cache_length
+                    max_cache_length is not None
+                    and attention_mask is not None
+                    and cache_length + input_ids.shape[1] > max_cache_length
             ):
                 attention_mask = attention_mask[:, -max_cache_length:]
 
@@ -1807,17 +1794,35 @@ class DeepseekV2ForCausalLM(DeepseekV2PreTrainedModel):
             position_ids = attention_mask.long().cumsum(-1) - 1
             position_ids.masked_fill_(attention_mask == 0, 1)
             if past_key_values:
-                position_ids = position_ids[:, -input_ids.shape[1] :]
+                position_ids = position_ids[:, -input_ids.shape[1]:]
+
+        if self.generation_config.cache_implementation == "static":
+            # generation with static cache
+            cache_position = kwargs.get("cache_position", None)
+            if cache_position is None:
+                past_length = 0
+            else:
+                past_length = cache_position[-1] + 1
+            input_ids = input_ids[:, past_length:]
+            position_ids = position_ids[:, past_length:]
+
+        # TODO @gante we should only keep a `cache_position` in generate, and do +=1.
+        # same goes for position ids. Could also help with continued generation.
+        cache_position = torch.arange(past_length, past_length + position_ids.shape[-1], device=position_ids.device)
 
         # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
         if inputs_embeds is not None and past_key_values is None:
             model_inputs = {"inputs_embeds": inputs_embeds}
         else:
-            model_inputs = {"input_ids": input_ids}
+            # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
+            # recompiles graphs as the stride of the inputs is a guard. Ref: https://github.com/huggingface/transformers/pull/29114
+            # TODO: use `next_tokens` directly instead.
+            model_inputs = {"input_ids": input_ids.contiguous()}
 
         model_inputs.update(
             {
-                "position_ids": position_ids,
+                "position_ids": position_ids.contiguous(),
+                "cache_position": cache_position,
                 "past_key_values": past_key_values,
                 "use_cache": kwargs.get("use_cache"),
                 "attention_mask": attention_mask,
@@ -1871,17 +1876,17 @@ class DeepseekV2ForSequenceClassification(DeepseekV2PreTrainedModel):
 
     @add_start_docstrings_to_model_forward(DeepseekV2_INPUTS_DOCSTRING)
     def forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[List[torch.FloatTensor]] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
+            self,
+            input_ids: torch.LongTensor = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[List[torch.FloatTensor]] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
     ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
         r"""
         labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
@@ -1921,7 +1926,7 @@ class DeepseekV2ForSequenceClassification(DeepseekV2PreTrainedModel):
         else:
             if input_ids is not None:
                 sequence_lengths = (
-                    torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
+                        torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
                 ).to(logits.device)
             else:
                 sequence_lengths = -1
@@ -1937,7 +1942,7 @@ class DeepseekV2ForSequenceClassification(DeepseekV2PreTrainedModel):
                 if self.num_labels == 1:
                     self.config.problem_type = "regression"
                 elif self.num_labels > 1 and (
-                    labels.dtype == torch.long or labels.dtype == torch.int
+                        labels.dtype == torch.long or labels.dtype == torch.int
                 ):
                     self.config.problem_type = "single_label_classification"
                 else:

deepseek_vl2/models/modeling_deepseek_vl_v2.py

@@ -1,4 +1,8 @@
 from attrdict import AttrDict
+from dataclasses import dataclass
+import logging
+import gc
+
 from einops import rearrange, repeat
 from typing import Optional, List, Tuple, Callable, Union
 
@@ -6,19 +10,27 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from transformers.utils import (
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+)
+from transformers.modeling_outputs import ModelOutput
 from transformers.configuration_utils import PretrainedConfig
 from transformers import (
     AutoConfig,
     AutoModelForCausalLM,
-    PreTrainedModel, GenerationConfig, LogitsProcessorList, StoppingCriteriaList,
+    PreTrainedModel
 )
-from transformers.generation.utils import GenerateOutput
+from transformers.utils import logging
 
 from .siglip_vit import VisionTransformer
 from .configuration_deepseek import DeepseekV2Config
 from .modeling_deepseek import DeepseekV2ForCausalLM
 
 
+logger = logging.get_logger(__name__)
+
+
 class MlpProjector(nn.Module):
 
     def __init__(self, cfg):
@@ -181,6 +193,45 @@ class MlpProjectorConfig(PretrainedConfig):
         super().__init__(**kwargs)
 
 
+@dataclass
+class DeepSeekVLV2CausalLMOutputWithPast(ModelOutput):
+    """
+    Base class for DeepSeek-VL2 causal language model (or autoregressive) outputs.
+
+    Args:
+        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+            Language modeling loss (for next-token prediction).
+        logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+        past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+            Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
+            `(batch_size, num_heads, sequence_length, embed_size_per_head)`)
+
+            Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+            `past_key_values` input) to speed up sequential decoding.
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+        rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+            The rope index difference between sequence length and multimodal rope.
+    """
+
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    past_key_values: Optional[List[torch.FloatTensor]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+    rope_deltas: Optional[torch.LongTensor] = None
+
+
 class DeepseekVLV2Config(PretrainedConfig):
     model_type = "deepseek_vl_v2"
     vision_config: VisionEncoderConfig
@@ -229,6 +280,8 @@ class DeepseekVLV2ForCausalLM(DeepseekVLV2PreTrainedModel):
     def __init__(self, config: DeepseekVLV2Config):
         super().__init__(config)
 
+        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+
         # ----------- vision encoder ------------
         vision_config = config.vision_config
         self.vision = VisionTransformer(
@@ -283,8 +336,8 @@ class DeepseekVLV2ForCausalLM(DeepseekVLV2PreTrainedModel):
     def prepare_inputs_embeds(
             self,
             input_ids: torch.LongTensor,
-            images: torch.FloatTensor,
-            images_seq_mask: torch.LongTensor,
+            images: Optional[torch.FloatTensor] = None,
+            images_seq_mask: Optional[torch.LongTensor] = None,
             images_spatial_crop: Optional[torch.LongTensor] = None,
             **ignore_kwargs
     ):
@@ -423,48 +476,222 @@ class DeepseekVLV2ForCausalLM(DeepseekVLV2PreTrainedModel):
 
         return input_embeds
 
-    def generate(
+    @torch.no_grad()
+    def incremental_prefilling(
             self,
-            inputs: Optional[torch.Tensor] = None,
-            generation_config: Optional[GenerationConfig] = None,
-            logits_processor: Optional[LogitsProcessorList] = None,
-            stopping_criteria: Optional[StoppingCriteriaList] = None,
-            prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
-            synced_gpus: Optional[bool] = None,
-            assistant_model: Optional["PreTrainedModel"] = None,
-            streamer: Optional["BaseStreamer"] = None,
-            negative_prompt_ids: Optional[torch.Tensor] = None,
-            negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+            input_ids: Optional[torch.LongTensor] = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+
+            images: Optional[torch.FloatTensor] = None,
+            images_seq_mask: Optional[torch.LongTensor] = None,
+            images_spatial_crop: Optional[torch.LongTensor] = None,
+            chunk_size: int = 1024
+    ):
+        if inputs_embeds is None:
+            inputs_embeds = self.prepare_inputs_embeds(
+                input_ids=input_ids,
+                images=images,
+                images_seq_mask=images_seq_mask,
+                images_spatial_crop=images_spatial_crop,
+            )
+
+            del images
+            del images_seq_mask
+            del images_spatial_crop
+
+            if attention_mask is not None:
+                attention_mask = attention_mask.to(inputs_embeds.device)
+
+            self._clear_cuda_cache()
+
+        bzs, seq_len, _ = inputs_embeds.shape
+        past_key_values = None
+
+        # remain the last token for the next forward
+        prefilling_len = seq_len - 1
+        for i in range(0, prefilling_len, chunk_size):
+            chunk_start = i
+            chunk_end = min(i + chunk_size, prefilling_len)
+            chunk_inputs_embeds = inputs_embeds[:, chunk_start: chunk_end]
+            chunk_attention_mask = attention_mask[:, 0: chunk_end]
+            # print(f"start = {chunk_start}, end = {chunk_end}, prefilling_len = {prefilling_len}, seq_len = {seq_len}")
+
+            # compute position_ids
+            if past_key_values is not None:
+                position_ids = torch.arange(
+                    chunk_start,
+                    chunk_end,
+                    dtype=torch.long,
+                    device=inputs_embeds.device
+                ).unsqueeze(0)
+                past_key_values = self._move_past_key_values_to_gpu(past_key_values, inputs_embeds.device)
+            else:
+                position_ids = None
+
+            # chunk-forward
+            with torch.no_grad():
+                outputs = self.forward(
+                    inputs_embeds=chunk_inputs_embeds,
+                    attention_mask=chunk_attention_mask,
+                    past_key_values=past_key_values,
+                    position_ids=position_ids,
+                    use_cache=True,
+                )
+                # update past_key_values
+                past_key_values = outputs.past_key_values
+                past_key_values = self._move_past_key_values_to_cpu(past_key_values)
+
+                del outputs, position_ids
+                self._clear_cuda_cache()
+
+        prefilling_key_values = []
+        for layer_past in past_key_values:
+            prefilling_key_values.append(
+                (
+                    layer_past[0][:, :, 0: prefilling_len, ...].to(inputs_embeds.device),
+                    layer_past[1][:, :, 0: prefilling_len, ...].to(inputs_embeds.device),
+                )
+            )
+
+        return inputs_embeds, prefilling_key_values
+
+    def forward(
+            self,
+            input_ids: Optional[torch.LongTensor] = None,
+
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[List[torch.FloatTensor]] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+
+            images: Optional[torch.FloatTensor] = None,
+            images_seq_mask: Optional[torch.LongTensor] = None,
+            images_spatial_crop: Optional[torch.LongTensor] = None,
+
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+            cache_position: Optional[torch.LongTensor] = None,
+    ):
+
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+        if inputs_embeds is None:
+            inputs_embeds = self.prepare_inputs_embeds(
+                input_ids=input_ids,
+                images=images,
+                images_seq_mask=images_seq_mask,
+                images_spatial_crop=images_spatial_crop,
+            )
+
+            if attention_mask is not None:
+                attention_mask = attention_mask.to(inputs_embeds.device)
+
+        # print(inputs_embeds.shape)
+        outputs = self.language.forward(
+            input_ids=None,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            labels=labels,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position
+        )
+
+        self._clear_cuda_cache()
+
+        return outputs
+
+    def _clear_cuda_cache(self):
+        """clear CUDA memory cache"""
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()
+
+    def _move_past_key_values_to_cpu(self, past_key_values):
+        # print(f"past_key_values -> cpu")
+        if past_key_values is None:
+            return None
+        return tuple(tuple(t.cpu() for t in layer) for layer in past_key_values)
+
+    def _move_past_key_values_to_gpu(self, past_key_values, device="cuda:0"):
+        # print(f"past_key_values -> gpu")
+        if past_key_values is None:
+            return None
+        return tuple(tuple(t.to(device) for t in layer) for layer in past_key_values)
+
+    def prepare_inputs_for_generation(
+            self,
+            input_ids,
+            past_key_values=None,
+            inputs_embeds=None,
+
+            images: Optional[torch.FloatTensor] = None,
+            images_seq_mask: Optional[torch.LongTensor] = None,
+            images_spatial_crop: Optional[torch.LongTensor] = None,
+
+            attention_mask=None,
+            cache_position=None,
+
+            pixel_values=None,
+            image_sizes=None,
+            num_logits_to_keep=None,
             **kwargs,
-    ) -> Union[GenerateOutput, torch.LongTensor]:
-        r"""
-        Generates sequences for models with a language modeling head. The method currently supports greedy decoding,
-        beam-search decoding, sampling with temperature, sampling with top-k or nucleus sampling. Beam-search decoding
-        is controlled by the `num_beams` parameter and the `num_return_sequences` parameter.
-
-        Parameters:
-            - `inputs` (optional) -- `torch.LongTensor` of shape `(batch, sequence_length)`:
-                The sequence used as a prompt for the generation. If `None`, generate for the model's prompt.
-            - `generation_config` (optional) -- `GenerationConfig`:
-                The generation config of the model.
-            - `logits_processor` (optional) -- `LogitsProcessorList`:
-                A list of instances of :class:`~transform
-        """
-
-        return self.language.generate(
-            inputs=inputs,
-            generation_config=generation_config,
-            logits_processor=logits_processor,
-            stopping_criteria=stopping_criteria,
-            prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
-            synced_gpus=synced_gpus,
-            assistant_model=assistant_model,
-            streamer=streamer,
-            negative_prompt_ids=negative_prompt_ids,
-            negative_prompt_attention_mask=negative_prompt_attention_mask,
+    ):
+        # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+        model_inputs = self.language.prepare_inputs_for_generation(
+            input_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            num_logits_to_keep=num_logits_to_keep,
             **kwargs,
         )
 
+        # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore
+        # Otherwise we need pixel values to be passed to model
+        cache_position = model_inputs["cache_position"]
+        if cache_position[0] == 0:
+            model_inputs["images"] = images
+            model_inputs["images_seq_mask"] = images_seq_mask
+            model_inputs["images_spatial_crop"] = images_spatial_crop
+
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (
+                tuple(
+                    past_state.index_select(0, beam_idx.to(past_state.device))
+                    for past_state in layer_past
+                ),
+            )
+        return reordered_past
+
 
 AutoConfig.register("vision", VisionEncoderConfig)
 AutoConfig.register("mlp_projector", MlpProjectorConfig)

deepseek_vl2/models/processing_deepseek_vl_v2.py

@@ -559,7 +559,7 @@ class DeepseekVLV2Processor(ProcessorMixin):
                 for j in range(0, best_width, self.image_size):
                     images_list.append(
                         self.image_transform(local_view.crop((j, i, j + self.image_size, i + self.image_size))))
-            
+
             """record height / width crop num"""
             num_width_tiles, num_height_tiles = best_width // self.image_size, best_height // self.image_size
             images_spatial_crop.append([num_width_tiles, num_height_tiles])