[feat] Add kv cache for InternVLA-N1 realworld deployment

Author: yuqiang-yang

internnav/agent/internvla_n1_agent_realworld.py

@@ -202,10 +202,10 @@ def step_s2(self, rgb, depth, pose, instruction, intrinsic, look_down = False):
                 **inputs, 
                 max_new_tokens=128, 
                 do_sample=False,
-                # use_cache=True,
-                # past_key_values=self.past_key_values,
+                use_cache=True,
+                past_key_values=self.past_key_values,
                 return_dict_in_generate=True,
-                # raw_input_ids=copy.deepcopy(inputs.input_ids),
+                raw_input_ids=copy.deepcopy(inputs.input_ids),
             )
         output_ids = outputs.sequences
 

internnav/model/basemodel/internvla_n1/internvla_n1.py

@@ -101,6 +101,49 @@ def __init__(self, config):
     def get_model(self):
         return self.model
 
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        cache_position=None,
+        position_ids=None,
+        use_cache=True,
+        pixel_values=None,
+        pixel_values_videos=None,
+        image_grid_thw=None,
+        video_grid_thw=None,
+        second_per_grid_ts=None,
+        **kwargs,
+    ):
+        # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+
+        model_inputs = super().prepare_inputs_for_generation(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            cache_position=cache_position,
+            position_ids=position_ids,
+            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,
+            use_cache=use_cache,
+            **kwargs,
+        )
+        # Qwen2-5-VL position_ids are prepareed with rope_deltas in forward
+        model_inputs["position_ids"] = None
+
+        # add for QwenVL kv cache
+        model_inputs["pixel_values"] = pixel_values
+        model_inputs["pixel_values_videos"] = pixel_values_videos
+
+        return model_inputs
+    
+    
     def forward(
         self,
         input_ids: Optional[torch.LongTensor] = None,
@@ -121,6 +164,7 @@ def forward(
         rope_deltas: Optional[torch.LongTensor] = None,
         cache_position: Optional[torch.LongTensor] = None,
         second_per_grid_ts: Optional[torch.Tensor] = None,
+        raw_input_ids: Optional[torch.LongTensor] = None,
     ) -> Union[Tuple, CausalLMOutputWithPast]:
         r"""
             labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
@@ -169,10 +213,11 @@ def forward(
 
         if inputs_embeds is None:
             inputs_embeds = self.model.embed_tokens(input_ids)
-            if pixel_values is not None:
+            n_image_tokens = (input_ids == self.config.image_token_id).sum().item()
+            if pixel_values is not None and n_image_tokens > 0:
                 pixel_values = pixel_values.type(self.visual.dtype)
                 image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
-                n_image_tokens = (input_ids == self.config.image_token_id).sum().item()
+                image_embeds = image_embeds[-n_image_tokens:]
                 n_image_features = image_embeds.shape[0]
                 if n_image_tokens != n_image_features:
                     raise ValueError(
@@ -232,6 +277,22 @@ def forward(
                     attention_mask,
                 )
                 self.rope_deltas = rope_deltas
+            elif n_image_tokens > 0: # using only for kv cache
+                attention_mask = attention_mask[:, :raw_input_ids.shape[1]]
+                position_ids, rope_deltas = self.get_rope_index(
+                    raw_input_ids,
+                    image_grid_thw,
+                    video_grid_thw,
+                    second_per_grid_ts,
+                    attention_mask,
+                )
+                delta = (
+                    (cache_position[0] + self.rope_deltas).to(inputs_embeds.device)
+                    if cache_position is not None
+                    else 0
+                )
+                position_ids = position_ids[:, :,-input_ids.shape[1]:]
+                self.rope_deltas = rope_deltas
             # then use the prev pre-calculated rope-deltas to get the correct position ids
             else:
                 batch_size, seq_length, _ = inputs_embeds.shape