Support Qwen 3VL in AdaScale ONNX (#6990)

Signed-off-by: Ashvin Kumar <quic_ashvkuma@quicinc.com>
Co-authored-by: Ashvin Kumar <quic_ashvkuma@quicinc.com>

Author: 2 people

GenAILab/qai_hub_lm/models/generator.py

@@ -743,11 +743,10 @@ def prefill(
             else [t for layer_kv in past_key_values for t in (layer_kv[0], layer_kv[1])]
         }
 
-        selected_seq_len = self._select_sequence_length(input_tokens.shape[1])
         slices_iter = self.slice_inputs_for_inference(
             input_tokens,
             attention_mask,
-            selected_seq_len,
+            self.sequence_length,  # always use max sequence length for slicing in prefill
             position_ids,
             **kwargs,
         )
@@ -767,7 +766,7 @@ def prefill(
                 input_ids=input_slice if input_ids is not None else None,
                 attention_mask=attention_mask_slice,
                 past_key_values=preconsumed_outputs["past_key_values"],
-                sequence_length=selected_seq_len,
+                sequence_length=self.sequence_length,
                 context_length=self.context_length,
                 pad_token=getattr(self.tokenizer, "eos_token_id", 0),
                 attention_mask_min=self.attention_mask_min,
@@ -798,7 +797,7 @@ def prefill(
             input_ids=input_slice if input_ids is not None else None,
             attention_mask=attention_mask_slice,
             past_key_values=preconsumed_outputs["past_key_values"],
-            sequence_length=selected_seq_len,
+            sequence_length=self.sequence_length,
             context_length=self.context_length,
             pad_token=getattr(self.tokenizer, "eos_token_id", 0),
             attention_mask_min=self.attention_mask_min,

GenAILab/qai_hub_lm/models/qwen2_vl.py

@@ -102,11 +102,15 @@ def generate_position_ids(
         self,
         *args,
         input_ids: torch.Tensor,
-        attention_mask: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
         **kwargs,
     ):
         num_new_tokens = input_ids.shape[1]
-        attention_mask = attention_mask[:, -num_new_tokens:]
+        attention_mask = (
+            attention_mask[:, -num_new_tokens:]
+            if attention_mask is not None
+            else torch.ones_like(input_ids, dtype=torch.int32)
+        )
 
         ctx = PositionIdContext(self.config, modeling_qwen2_5_vl.Qwen2_5_VLModel)
         position_ids, *_ = modeling_qwen2_5_vl.Qwen2_5_VLModel.get_rope_index(

GenAILab/qai_hub_lm/models/qwen3_vl.py

@@ -165,11 +165,15 @@ def generate_position_ids(
         self,
         *args,
         input_ids: torch.Tensor,
-        attention_mask: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
         **kwargs,
     ):
         num_new_tokens = input_ids.shape[1]
-        attention_mask = attention_mask[:, -num_new_tokens:]
+        attention_mask = (
+            attention_mask[:, -num_new_tokens:]
+            if attention_mask is not None
+            else torch.ones_like(input_ids, dtype=torch.int32)
+        )
 
         ctx = PositionIdContext(self.config, modeling_qwen3_vl.Qwen3VLModel)
         position_ids, *_ = modeling_qwen3_vl.Qwen3VLModel.get_rope_index(

GenAILab/tests/unit/models/test_generator.py

@@ -948,11 +948,11 @@ def test_prefill_uses_large_seq_len(self, multi_seq_gen):
         assert isinstance(slices[0], OrderedDict)
         assert slices[0]["input_ids"].shape[1] == 8
 
-    def test_prefill_single_token_uses_small_seq_len(self, multi_seq_gen):
+    def test_prefill_single_token_uses_max_seq_len(self, multi_seq_gen):
         input_ids = torch.randint(0, 256, (1, 1))
         slices = list(multi_seq_gen.prefill(input_ids=input_ids))
         assert len(slices) == 1
-        assert slices[0]["input_ids"].shape[1] == 1
+        assert slices[0]["input_ids"].shape[1] == max(multi_seq_gen.sequence_lengths)
 
     def test_three_sequence_lengths(self, model, tokenizer):
         gen = Generator(

TrainingExtensions/onnx/src/python/aimet_onnx/experimental/adascale/adascale_optimizer.py

@@ -73,6 +73,9 @@ class AdaScaleModelConfig:
     "qwen3": AdaScaleModelConfig(
         model_type="qwen3", beta_gamma_lr=1e-3, scales_lr=5e-4
     ),
+    "qwen3_vl": AdaScaleModelConfig(
+        model_type="qwen3", beta_gamma_lr=1e-3, scales_lr=5e-4
+    ),
     "phi3": AdaScaleModelConfig(model_type="phi3", beta_gamma_lr=1e-3, scales_lr=5e-4),
     "qwen2_5_vl": AdaScaleModelConfig(
         model_type="qwen2_5_vl", beta_gamma_lr=1e-3, scales_lr=5e-4
@@ -148,12 +151,15 @@ def apply_adascale(
                 )
 
             # create a list of common input names to be used for graph slicing and populating input_list
+            # Exclude primary sequence inputs (consumed upstream by embedding layer)
+            # and past_key_*/past_value_* (handled per-block below)
             common_input_names = []
             for name in graph_input_names:
-                if "attention" in name:
-                    common_input_names.append(name)
-                if "position" in name:
-                    common_input_names.append(name)
+                if name == "inputs_embeds" or "input_ids" in name:
+                    continue
+                if "past_key" in name or "past_value" in name:
+                    continue
+                common_input_names.append(name)
 
             del sim.session
             gc.collect()
@@ -198,7 +204,7 @@ def apply_adascale(
                         ):
                             block_kv_tensor_names.append(name)
 
-                    block_input_names = common_input_names
+                    block_input_names = list(common_input_names)
                     if len(block_kv_tensor_names) > 0:
                         if len(block_kv_tensor_names) != 2:
                             raise RuntimeError(