Adding SingleQPC

Signed-off-by: Mohit Soni <[email protected]>

Author: mohiso22

QEfficient/base/pytorch_transforms.py

@@ -129,11 +129,11 @@ class SplitGateUpWeightsTransform(PytorchTransform):
     def apply(cls, model: nn.Module) -> Tuple[nn.Module, bool]:
         transformed = False
 
-        model = model.language_model if hasattr(model, "language_model") else model
+        model_tmp = model.language_model if hasattr(model, "language_model") else model
 
-        num_layers = len(model.model.layers)
+        num_layers = len(model_tmp.model.layers)
         delete_fused_key = True
-        sd = model.state_dict()
+        sd = model_tmp.state_dict()
         for layer_idx in range(num_layers):
             # ---- build the textual prefix once per layer ----------
             prefix = f"model.layers.{layer_idx}.feed_forward.experts."
@@ -148,7 +148,7 @@ def apply(cls, model: nn.Module) -> Tuple[nn.Module, bool]:
             ffn_dim = two_I // 2
             gate, up = fused.split(ffn_dim, dim=-1)  # views – no copy
 
-            experts = model.model.layers[layer_idx].feed_forward.experts
+            experts = model_tmp.model.layers[layer_idx].feed_forward.experts
             experts.gate_proj.data.copy_(gate)
             experts.up_proj.data.copy_(up)
 
@@ -161,6 +161,8 @@ def apply(cls, model: nn.Module) -> Tuple[nn.Module, bool]:
 
             logger.info(f"[layer {layer_idx:02d}] loaded gate_proj & up_proj from fused tensor  (shape {fused.shape})")
             transformed = True
+
+        model.language_model = model_tmp
         return model, transformed
 
 

QEfficient/transformers/models/llama4/modeling_llama4.py

@@ -883,6 +883,32 @@ def get_qeff_vision_encoder(self):
     def get_qeff_language_decoder(self):
         return QEffLlama4DecoderWrapper(self)
 
+    def forward(self, input_ids, position_ids, pixel_values, index, past_key_values):
+        inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
+        vision_feature_layer = self.config.vision_config.vision_feature_layer
+        vision_feature_select_strategy = self.config.vision_config.vision_feature_select_strategy
+        image_features = self.get_image_features(
+            pixel_values=pixel_values,
+            vision_feature_layer=vision_feature_layer,
+            vision_feature_select_strategy=vision_feature_select_strategy,
+            image_sizes=None,
+        )
+        vision_flat = image_features.view(-1, image_features.size(-1))
+        projected_vision_flat = self.multi_modal_projector(vision_flat)
+        selected = input_ids == self.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 = projected_vision_flat.unsqueeze(0)[indices0, indices1]
+        image_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_embeds)
+        outputs = self.language_model(
+            inputs_embeds=inputs_embeds, position_ids=position_ids, past_key_values=past_key_values, use_cache=True
+        )
+        next_index = (indices1.max() + 1).unsqueeze(0).unsqueeze(0)
+        index = torch.where(index < next_index, next_index, index)
+        return outputs.logits, pixel_values, index, outputs.past_key_values
+
     def get_specializations(
         self,
         batch_size: int,
@@ -963,6 +989,7 @@ def get_onnx_dynamic_axes(self, kv_offload: bool = False):
             dynamic_axes["vision"] = vision_dynamic_axes
             dynamic_axes["lang"] = lang_dynamic_axes
         else:
+            lang_dynamic_axes.pop("vision_embeds")
             dynamic_axes = {**vision_dynamic_axes, **lang_dynamic_axes}
         return dynamic_axes
 
@@ -981,6 +1008,7 @@ def get_output_names(self, kv_offload: bool = False):
             output_names["lang"] = lang_output_names
         else:
             lang_output_names.insert(1, "pixel_values_RetainedState")
+            lang_output_names.insert(2, "index_output")
             return lang_output_names
         return output_names
 

QEfficient/transformers/models/modeling_auto.py

@@ -1068,8 +1068,6 @@ def cloud_ai_100_generate(
         generated_ids = np.full((batch_size, generation_len + 1), pad_token_id)
 
         # Prepare inputs for prefill
-        prefill_start = perf_counter()
-
         inputs["input_ids"] = torch.nn.functional.pad(
             inputs["input_ids"],
             (0, padded_len - input_ids_length),
@@ -1090,16 +1088,18 @@ def cloud_ai_100_generate(
             inputs["pixel_values"] = inputs["pixel_values"].astype("float16")
 
         inputs["position_ids"] = np.where(inputs.pop("attention_mask"), np.arange(padded_len), -1)
+        inputs["index"] = np.array([[0]])
 
         qpc_session.activate()
+        chunk_inputs = inputs.copy()
+        prefill_start = perf_counter()
 
         # Run prefill
-
         for i in range(num_chunks):
-            chunk_inputs = inputs.copy()
             chunk_inputs["input_ids"] = inputs["input_ids"][:, i * prefill_seq_len : (i + 1) * prefill_seq_len]
             chunk_inputs["position_ids"] = inputs["position_ids"][:, i * prefill_seq_len : (i + 1) * prefill_seq_len]
             outputs = qpc_session.run(chunk_inputs)
+            chunk_inputs["index"] = outputs["index_output"]
 
         prefill_time = perf_counter() - prefill_start
         # Get first token