Updating Wrappers for Merging and Chunking in DecoderWrapper (#404)

Signed-off-by: Mohit Soni <quic_mohisoni@quicinc.com>
Signed-off-by: Mohit Soni <mohisoni@qti.qualcomm.com>

Author: mohiso22

QEfficient/transformers/models/gemma3/modeling_gemma3.py

@@ -560,16 +560,9 @@ def __init__(self, model):
         self.model = model
         self.model.vision_model = self.model.vision_tower
 
-    def forward(self, input_ids, pixel_values):
-        inputs_embeds = self.model.get_input_embeddings()(input_ids)
-        B, N, C = inputs_embeds.shape
+    def forward(self, pixel_values):
         image_features = self.model.get_image_features(pixel_values=pixel_values)
-        selected = input_ids == self.model.config.image_token_index
-        indices1 = selected.to(torch.int64).cumsum(1) - 1
-        indices0 = torch.arange(selected.unsqueeze(0).shape[0]).view(-1, 1)
-        image_features_expanded = image_features.reshape(-1, C).unsqueeze(0)[indices0, indices1]
-        image_input_embeds = torch.where(selected.unsqueeze(-1), image_features_expanded, inputs_embeds)
-        return image_input_embeds
+        return image_features
 
 
 class QEffGemma3DecoderWrapper(nn.Module):
@@ -579,14 +572,21 @@ def __init__(self, model):
         self.language_model = self.model.language_model
         self.config = self.model.config
 
-    def forward(self, input_ids, vision_embeds, position_ids, past_key_values):
-        image_embeds = vision_embeds[:, : input_ids.shape[1], :]
-        inputs_embeds = self.model.language_model.get_input_embeddings()(input_ids)
-        inputs_embeds = torch.where(input_ids.shape[1] == torch.tensor(1), inputs_embeds, image_embeds)
+    def forward(self, input_ids, vision_embeds, position_ids, index, past_key_values):
+        inputs_embeds = self.model.get_input_embeddings()(input_ids)
+        B, N, C = inputs_embeds.shape
+        selected = input_ids == self.model.config.image_token_index
+        indices1 = selected.to(torch.int64).cumsum(1) - 1
+        indices1 = torch.where(indices1 != -1, indices1 + index, indices1)
+        indices0 = torch.arange(selected.unsqueeze(0).shape[0]).view(-1, 1)
+        image_features_expanded = vision_embeds.reshape(-1, C).unsqueeze(0)[indices0, indices1]
+        image_input_embeds = torch.where(selected.unsqueeze(-1), image_features_expanded, inputs_embeds)
+        inputs_embeds = torch.where(input_ids.shape[1] == torch.tensor(1), inputs_embeds, image_input_embeds)
         outputs = self.model.language_model(
             inputs_embeds=inputs_embeds, position_ids=position_ids, past_key_values=past_key_values, use_cache=True
         )
-        return outputs.logits, vision_embeds, outputs.past_key_values
+        index = (indices1.max() + 1).unsqueeze(0).unsqueeze(0)
+        return outputs.logits, vision_embeds, index, outputs.past_key_values
 
 
 class QEffGemma3ForConditionalGeneration(Gemma3ForConditionalGeneration):
@@ -605,24 +605,20 @@ def get_specializations(
         kv_offload: bool = False,
         **compiler_options,
     ):
-        vision_seq_len = compiler_options.pop("vision_seq_len", None)
-        if vision_seq_len is None:
-            # TODO: Check properly for Gemma3, Not verified yet.
-            vision_seq_len = 512  # for Gemma3 Vision feature shape is (1, 4096, 1152) --> 1152 is hidden size)
-
         prefill_seq_len = prefill_seq_len if prefill_seq_len else 32
         ctx_len = ctx_len if ctx_len else constants.INTERN_CTX_LEN
         if img_size is None and hasattr(self.config.vision_config, "image_size"):
             img_size = getattr(self.config.vision_config, "image_size")
         elif img_size is None:
             img_size = 896  # FIXME based on gemma3 Image size
             logger.warning("Setting img_size to be 336, as it was neither passed nor found in vision_config")
+        mm_tokens_per_image = getattr(self.config, "mm_tokens_per_image", 256)
 
         vision = [
             {
                 "batch_size": batch_size,
                 "img_size": img_size,
-                "seq_len": vision_seq_len,
+                "seq_len": prefill_seq_len,
                 "ctx_len": ctx_len,
             }
         ]
@@ -632,14 +628,14 @@ def get_specializations(
                 "seq_len": prefill_seq_len,
                 "ctx_len": ctx_len,
                 "img_size": img_size,
-                "chunk_length": prefill_seq_len,
+                "mm_tokens_per_image": mm_tokens_per_image,
             },
             {
                 "batch_size": batch_size,
                 "seq_len": "1",
                 "ctx_len": ctx_len,
                 "img_size": img_size,
-                "chunk_length": prefill_seq_len,
+                "mm_tokens_per_image": mm_tokens_per_image,
             },
         ]
 
@@ -658,9 +654,8 @@ def get_onnx_dynamic_axes(self, kv_offload: bool = False):
         lang_dynamic_axes = {}
         lang_dynamic_axes["input_ids"] = {0: "batch_size", 1: "seq_len"}
         lang_dynamic_axes["position_ids"] = {0: "batch_size", 1: "seq_len"}
-        lang_dynamic_axes["vision_embeds"] = {0: "batch_size", 1: "chunk_length"}
+        lang_dynamic_axes["vision_embeds"] = {0: "batch_size", 1: "mm_tokens_per_image"}
         vision_dynamic_axes["pixel_values"] = {0: "batch_size", 2: "img_size", 3: "img_size"}
-        vision_dynamic_axes["input_ids"] = {0: "batch_size", 1: "seq_len"}
 
         pkv_dynamic_axes = {0: "batch_size", 2: "ctx_len"}
         for i in range(self.language_model.config.num_hidden_layers):
@@ -685,6 +680,7 @@ def get_output_names(self, kv_offload: bool = False):
         output_names = {}
         if kv_offload:
             lang_output_names.insert(1, "vision_embeds_RetainedState")
+            lang_output_names.insert(2, "index_output")
             output_names["vision"] = vision_output_names
             output_names["lang"] = lang_output_names
         else:
@@ -698,12 +694,13 @@ def get_dummy_inputs(self, kv_offload: bool = False):
         else:
             img_size = 896
 
+        mm_tokens_per_image = getattr(self.config, "mm_tokens_per_image", 256)
         # Define shapes
         inputs_shapes = {}
         inputs_shapes["input_ids"] = (constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE, constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN)
         inputs_shapes["vision_embeds"] = (
             1,  # constants.INTERN_NUM_PATCHES,
-            constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN,  # constants.INTERN_FEATURE_SIZE,
+            mm_tokens_per_image,  # constants.INTERN_FEATURE_SIZE,
             self.language_model.config.hidden_size,  # 5120
         )
         inputs_shapes["position_ids"] = (
@@ -716,20 +713,20 @@ def get_dummy_inputs(self, kv_offload: bool = False):
             img_size,
             img_size,
         )
+        inputs_shapes["index"] = (1, 1)
 
         # Define inputs
         vision_inputs = {}
         lang_inputs = {}
         vision_inputs["pixel_values"] = torch.zeros((inputs_shapes["pixel_values"]), dtype=torch.float32)
-        vision_inputs["input_ids"] = torch.zeros((inputs_shapes["input_ids"]), dtype=torch.int64)
         lang_inputs["input_ids"] = torch.zeros((inputs_shapes["input_ids"]), dtype=torch.int64)
         lang_inputs["vision_embeds"] = torch.zeros((inputs_shapes["vision_embeds"]), dtype=torch.float32)
         lang_inputs["position_ids"] = (
             torch.arange(constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN, dtype=torch.int64)
             .view(1, constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN)
             .repeat(constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE, 1)
         )
-
+        lang_inputs["index"] = torch.zeros((inputs_shapes["index"]), dtype=torch.int64)
         # Add data for KV
         kv_cache_shape = get_padding_shape_from_config(
             config=self.language_model.config,

QEfficient/transformers/models/modeling_auto.py

@@ -751,8 +751,8 @@ def kv_offload_generate(
         input_len = inputs["attention_mask"].sum(1, keepdims=True)
         input_ids_length = inputs["input_ids"].shape[1]
         num_chunks = -(input_ids_length // -prefill_seq_len)  # ceil divide without float
-        # padded_len = num_chunks * prefill_seq_len  # Convert to a multiple of prompt_len
-        padded_len = vision_session.bindings[vision_session.binding_index_map["input_ids"]].dims[1]
+        padded_len = num_chunks * prefill_seq_len  # Convert to a multiple of prompt_len
+
         if generation_len is None:
             generation_len = ctx_len - input_len.max()
         assert generation_len > 0, "generation length should be greater than zero"
@@ -783,39 +783,39 @@ def kv_offload_generate(
         }
 
         vision_inputs["pixel_values"] = vision_inputs["pixel_values"].astype("float16")
-        vision_inputs["input_ids"] = inputs["input_ids"]
         vision_start = perf_counter()
         vision_outputs = vision_session.run(vision_inputs)
         vision_end = perf_counter()
 
         lang_inputs = {k: v for k, v in inputs.items() if k not in vision_inputs}
-        lang_inputs["input_ids"] = inputs["input_ids"]
         lang_inputs["position_ids"] = np.where(
             lang_inputs.pop("attention_mask"), np.arange(padded_len), -1
         )  # Need to use -1 as position_ids for invalid tokens
 
         vision_session.deactivate()
         lang_session.activate()
         lang_inputs["vision_embeds"] = vision_outputs["vision_embeds"]
-        # lang_session.set_buffers(vision_outputs)
+        lang_session.set_buffers(vision_outputs)
         prefill_start = perf_counter()
         # Run prefill
+        chunk_inputs = lang_inputs.copy()
+        chunk_inputs["index"] = np.array([[0]])
         for i in range(num_chunks):
-            chunk_inputs = lang_inputs.copy()
             chunk_inputs["input_ids"] = lang_inputs["input_ids"][:, i * prefill_seq_len : (i + 1) * prefill_seq_len]
             chunk_inputs["position_ids"] = lang_inputs["position_ids"][
                 :, i * prefill_seq_len : (i + 1) * prefill_seq_len
             ]
-            chunk_inputs["vision_embeds"] = lang_inputs["vision_embeds"][
-                :, i * prefill_seq_len : (i + 1) * prefill_seq_len
-            ]
             outputs = lang_session.run(chunk_inputs)
+            chunk_inputs["index"] = outputs["index_output"]
 
         prefill_time = perf_counter() - prefill_start + vision_end - vision_start
-        lang_inputs["vision_embeds"] = lang_inputs["vision_embeds"][:, :prefill_seq_len]
         # Skip inputs/outputs again
         lang_session.skip_buffers(
-            [x for x in lang_session.input_names + lang_session.output_names if x.startswith("past_")]
+            [
+                x
+                for x in lang_session.input_names + lang_session.output_names
+                if x.startswith("past_") or x.endswith("_RetainedState")
+            ]
         )
 
         # Get first token