Qualcomm AI Engine Direct - Decouple quantization and compile graphs for faster VLM/LLM PTQ (pytorch#19220)

### Summary
 - Calibrate decoder using prefill stage only (full chunk tokens)
 - Remove need for AR-N calibration loops
 - Significantly reduce calibration overhead

| model name | before</br> Time(sec) | after</br> Time(sec) | speedup |
|------------|--------------|-------------|---------|
|    gemma-2b        |      1216     |    259    |  4.69x  |
|      gemma2-2b      |      1827      |   382   |    4.78x     |
|    gemma3-1b        |       907     |    218   |     4.16x    |
|      glm-1_5b     |        963       |    167   |    5.76x     |
| granite_3_3-2b | 1545 | 304 | 5.08x |
|      llama3_2-1b      |       1237    |      285  |    4.34×     |
|   llama3_2-3b        |      2286     |   813   |     2.81x    |
|      phi_4_mini      |       2824     |   363   |   7.77x     |
|    qwen2_5-0_5b        |     486      |     119       |   4.08x      |
|     qwen2_5-1_5b      |      1068   |     220   |   4.86×    |
| qwen3-0_6b | 1013 | 158 | 6.41× |
|     qwen3-1_7b      |       1478    |     283  |  5.22×      |
|     smollm2_135m      |       399      |    122         |  3.27×   |
|     smollm3-3b      |      2065      |   431   | 4.79x  |
| smolvlm_500m_instruct | 170 | 131 | 1.30× |
| internvl3_1b | 170 | 103 | 1.65x |
| granite_speech_3_3-2b | 447 | 215 | 2.07x |


This change decouples the quantization graph from the graph used for
subsequent lowering, so calibration no longer depends on the AR-N
decoding flow.

Previously, we were running calibration directly on the graph shaped for
lowering (with fixed AR-N constraints). That forced us into an
autoregressive loop (AR1 per step), which was both inefficient and slow
since we never saw the full sequence context in a single pass.

With this update, calibration is done once during the prefill stage
using the full tokens chunk. This gives us much better coverage in a
single run and completely removes the need for iterative decoding during
calibration.

After quantization, we take the KV cache encodings from the output,
override the input KV cache encodings, and then propagate those into the
graph that will later be lowered. This keeps everything consistent
without needing to recalibrate on that graph.

Result: same accuracy, significantly faster calibration, and a much
cleaner separation between quantization and lowering

### Test plan
Test CI in `TestExampleLLMScript` and `TestExampleMultimodalityScript`

Author: DannyYuyang-quic

backends/qualcomm/tests/test_qnn_delegate.py

@@ -7608,6 +7608,11 @@ def test_static_llm_model(self):  # noqa: C901
             case "sqnr":
                 cmds.extend(
                     [
+                        "--skip_user_prompt_calibration",
+                        "--tasks",
+                        "wikitext",
+                        "--limit",
+                        "1",
                         "--eval_methods",
                         "sqnr_eval",
                     ]

examples/qualcomm/oss_scripts/llama/decoder_runtime_evaluator.py

@@ -424,7 +424,7 @@ def __init__(
                 self.max_seq_length = pte_max_context_len
 
     def run(self, prompt):
-        golden_logits, _ = INFERENCE_REGISTRY[True](
+        result = INFERENCE_REGISTRY[True](
             get_example_inputs=self.get_example_inputs,
             prompt=prompt,
             module=self.source_model,
@@ -433,6 +433,7 @@ def run(self, prompt):
             use_i64_token=self.args.embedding_quantize is not None,
             collect_logits=True,
         )
+        golden_logits = result.logits
 
         input_file_name = f"{self.args.artifact}/input_tokens.raw"
 

examples/qualcomm/oss_scripts/llama/decoder_utils.py

@@ -77,6 +77,13 @@ class DecoderInputs:
     embedding: Optional[torch.Tensor] = None
 
 
+@dataclass
+class DecoderOutputs:
+    logits: Optional[torch.Tensor] = None
+    token_list: Optional[List[int]] = None
+    input_samples: Optional[List] = None
+
+
 class GraphModuleCalibrationWrapper(EagerEvalWrapper):
     """
     A wrapper class for calibration
@@ -94,6 +101,7 @@ def __init__(  # noqa: C901
         get_example_inputs: Callable,
         use_i64_token: bool,
         seq_mse_candidates: int,
+        collect_input_samples: bool = False,
     ):
         # n seq len = n-1 cache len, so we len(inps) = n-1 during _model_call
         assert max_seq_length is not None, "max_seq_length must be provided"
@@ -108,30 +116,32 @@ def __init__(  # noqa: C901
         self.use_i64_token = use_i64_token
         self.seq_mse_candidates = seq_mse_candidates
         self._input_samples = None
+        self.collect_input_samples = collect_input_samples
 
     def get_input_samples(self):
         return self._input_samples
 
     def _model_call(self, inps):
-        all_logits = None
         kwargs = {}
         if self._use_kv_cache:
             kwargs["ar_len"] = self.ar_len
             kwargs["seq_mse_candidates"] = self.seq_mse_candidates
 
-        all_logits, self._input_samples = INFERENCE_REGISTRY[self._use_kv_cache](
+        result = INFERENCE_REGISTRY[self._use_kv_cache](
             self.get_example_inputs,
             inps,
             self._model,
             self._tokenizer,
             max_seq_len=self.max_seq_length,
             use_i64_token=self.use_i64_token,
             collect_logits=True,
+            collect_input_samples=self.collect_input_samples,
             **kwargs,
         )
+        self._input_samples = result.input_samples
         # one shot is enough for seq mse
         self.seq_mse_candidates = 0
-        return all_logits
+        return result.logits
 
 
 class LookaheadDecoder:
@@ -727,7 +737,8 @@ def kv_inference(  # noqa: C901
     collect_logits=False,
     seq_mse_candidates=0,
     lookahead_config=None,
-):
+    collect_input_samples=False,
+) -> DecoderOutputs:
     input_samples = []  # Record input sample for quantization error analysis
     is_multimodal = all(
         [
@@ -814,6 +825,7 @@ def kv_inference(  # noqa: C901
 
     # record total input tokens and generated tokens
     total_token_list = prompt_token_list
+    last_token_in_prompt = prompt_token_list[-1] if len(prompt_token_list) > 0 else None
 
     # 3. prepare decoder inputs
     inputs = DecoderInputs(
@@ -841,28 +853,33 @@ def kv_inference(  # noqa: C901
 
         # Phase 2: Generate tokens until the EOS token is generated or max_seq_len is reached.
         # When run on wikitext for ppl evaluation, this while-loop is not expected to run.
-        generate_input_sample = _generate(
-            inputs,
-            cur_pos,
-            module,
-            tokenizer,
-            tok_embedding,
-            ar_len,
-            max_seq_len,
-            k_caches,
-            v_caches,
-            total_token_list,
-            lookahead_config,
-        )
-        if generate_input_sample is not None:
-            input_samples.append(generate_input_sample)
-        else:
-            input_samples.append(prefill_input_sample)
+        generate_input_sample = None
+        if last_token_in_prompt != tokenizer.eos_id:
+            generate_input_sample = _generate(
+                inputs,
+                cur_pos,
+                module,
+                tokenizer,
+                tok_embedding,
+                ar_len,
+                max_seq_len,
+                k_caches,
+                v_caches,
+                total_token_list,
+                lookahead_config,
+            )
+
+        if collect_input_samples:
+            input_samples.append(generate_input_sample or prefill_input_sample)
 
     logging.info(f"kv inference result:\n{tokenizer.decode(total_token_list)}")
     if collect_logits:
         result_logits = torch.cat(result_logits, dim=1)
-    return result_logits, input_samples
+    return DecoderOutputs(
+        logits=result_logits if collect_logits else None,
+        token_list=total_token_list,
+        input_samples=input_samples if collect_input_samples else None,
+    )
 
 
 @register_inference(use_kv_cache=False)
@@ -878,7 +895,8 @@ def prefill_inference(
     max_seq_len=512,
     use_i64_token=False,
     collect_logits=False,
-):
+    collect_input_samples=False,
+) -> DecoderOutputs:
     input_samples = None  # Record input sample for quantization error analysis
     is_multimodal = all(
         [
@@ -946,7 +964,11 @@ def prefill_inference(
             pos += 1
     if isinstance(prompt, str):
         logging.info(f"prefill inference result:\n{tokenizer.decode(token_list)}")
-    return result_logits, [input_samples]
+    return DecoderOutputs(
+        logits=result_logits if collect_logits else None,
+        token_list=token_list,
+        input_samples=[input_samples] if collect_input_samples else None,
+    )
 
 
 def graph_module_inference(
@@ -968,7 +990,8 @@ def graph_module_inference(
     event_name: Optional[str] = None,
     seq_mse_candidates: int = 0,
     lookahead_config: Optional[Tuple[int]] = None,
-):
+    collect_input_samples: bool = False,
+) -> DecoderOutputs:
     """
     This function supports model execution from static nn.Module decoder model
     all the way to edge program.
@@ -984,7 +1007,7 @@ def graph_module_inference(
             kwargs["ar_len"] = ar_len
             kwargs["lookahead_config"] = lookahead_config
 
-        _, input_samples = INFERENCE_REGISTRY[use_kv_cache](
+        result = INFERENCE_REGISTRY[use_kv_cache](
             get_example_inputs,
             prompt,
             module,
@@ -996,10 +1019,11 @@ def graph_module_inference(
             max_seq_len=max_seq_len,
             use_i64_token=use_i64_token,
             collect_logits=False,
+            collect_input_samples=collect_input_samples,
             **kwargs,
         )
         logging.info(f"Prompt summary for {event_name}")
-        return input_samples
+        return result
     else:
         calibration_wrapper = GraphModuleCalibrationWrapper(
             model=module,
@@ -1010,6 +1034,7 @@ def graph_module_inference(
             get_example_inputs=get_example_inputs,
             use_i64_token=use_i64_token,
             seq_mse_candidates=seq_mse_candidates,
+            collect_input_samples=collect_input_samples,
         )
         with torch.no_grad():
             eval_results = simple_evaluate(
@@ -1022,4 +1047,4 @@ def graph_module_inference(
         for task, res in eval_results["results"].items():
             logging.info(f"{task}: {res}")
 
-        return calibration_wrapper.get_input_samples()
+        return DecoderOutputs(input_samples=calibration_wrapper.get_input_samples())

examples/qualcomm/oss_scripts/llama/llama.py

@@ -577,6 +577,12 @@ def _build_parser():
         help="Enable automatic quant recipe suggestion in PTQ",
     )
 
+    parser.add_argument(
+        "--skip_user_prompt_calibration",
+        action="store_true",
+        help="Skip using user prompt for calibration. Useful when only dataset-based calibration is desired.",
+    )
+
     return parser
 
 
@@ -676,6 +682,9 @@ def export_llama(args) -> None:
     assert (
         not is_multimodal or args.use_attention_sink is None
     ), "Multimodal models currently do not support attention sink feature."
+    assert (
+        not is_multimodal or not args.skip_user_prompt_calibration
+    ), "--skip_user_prompt_calibration is not supported for multimodal models (VLM/ALM) as they do not support task-based calibration yet."
 
     if args.pre_gen_pte:
         text_decoder_pte_path = f"{args.pre_gen_pte}/{pte_filenames[TEXT_DECODER]}.pte"

examples/qualcomm/oss_scripts/llama/wrappers/base_component.py

@@ -40,6 +40,7 @@
 class Mode(Enum):
     PREFILL = 1
     DECODE = 2
+    CALIBRATE = 3
 
 
 def log_info(func):
@@ -83,7 +84,7 @@ def process_model_args(
         model_args: ModelArgs object to be modified.
         quant_recipe: Quantization recipe to be used.
         config: LLMModelConfig object to be used.
-        mode: Mode of operation (PREFILL or DECODE).
+        mode: Mode of operation (PREFILL, DECODE, or CALIBRATE).
     """
     # TODO: support batch inputs if necessary
     if mode == Mode.DECODE:
@@ -95,13 +96,19 @@ def process_model_args(
             if control_args.model_mode == "lookahead"
             else 1
         )
-    else:
+    elif mode == Mode.PREFILL:
         ar_len = control_args.prefill_ar_len
+    elif mode == Mode.CALIBRATE:
+        ar_len = control_args.max_context_len
+    else:
+        raise ValueError(f"Unsupported mode: {mode}")
 
     model_args.max_batch_size = 1
     model_args.max_seq_len = control_args.max_seq_len
     model_args.max_context_len = control_args.max_context_len
-    model_args.use_kv_cache = control_args.max_context_len != ar_len
+    model_args.use_kv_cache = (
+        control_args.max_context_len != ar_len or mode == Mode.CALIBRATE
+    )
     model_args.enable_r3 = config.r3
     model_args.ar_len = ar_len
     model_args.kv_io_bit_width = quant_recipe.get_kv_io_bit_width()