Qualcomm AI Engine Direct - [GenAI Pipeline] PR5: Model preparation & quantization strategy implementations

backends/qualcomm/genai_pipeline/configs/model_preparation_output_config.py

@@ -8,7 +8,7 @@
 
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Iterable, Optional, TYPE_CHECKING
+from typing import Any, Iterable, Optional, Tuple, TYPE_CHECKING
 
 if TYPE_CHECKING:
     from torch import nn
@@ -27,6 +27,12 @@ class ModelPreparationOutputConfig:
     Attributes:
         model_module: The prepared nn.Module ready for quantization.
         tokenizer: The tokenizer instance for encoding/decoding text.
+        example_inputs: Positional example inputs for ``torch.export``, derived
+            from the **model** (never from ``calibration_data``): they carry the
+            exported graph's signature, its zero-initialized KV caches, and the
+            AR length baked in because HTP has no dynamic shapes. The dependency
+            runs model -> dataset, not the reverse -- the calibration dataset's
+            attention-mask schema is itself derived from this tuple.
         calibration_data: Calibration samples. Any Iterable[Tuple[Tensor, ...]],
             including a DataLoader with a custom collate_fn.
         runtime_tokenizer_path: Path to the runtime tokenizer **file** (not the
@@ -36,6 +42,7 @@ class ModelPreparationOutputConfig:
 
     model_module: Optional["nn.Module"] = None
     tokenizer: Any = None
+    example_inputs: Optional[Tuple[Any, ...]] = None
     calibration_data: Optional[Iterable[Any]] = None
     runtime_tokenizer_path: Optional[Path] = None
     chat_template: Optional[str] = None

backends/qualcomm/genai_pipeline/configs/quantization_input_config.py

@@ -7,7 +7,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Any, Dict, Iterable, Optional, TYPE_CHECKING
+from typing import Any, Dict, Iterable, Optional, Tuple, TYPE_CHECKING
 
 if TYPE_CHECKING:
     from executorch.backends.qualcomm.serialization.qc_schema import (
@@ -21,10 +21,11 @@
 class QuantizationInputConfig:
     """Input configuration for the quantization stage.
 
-    ``model_module`` and ``calibration_data`` are ``Optional`` only because the
-    orchestrator builds this from the previous stage's output, which is empty
-    when model preparation is skipped. **Both are required once the quantization
-    stage executes**, and strategies should validate their presence.
+    ``model_module``, ``example_inputs`` and ``calibration_data`` are
+    ``Optional`` only because the orchestrator builds this from the previous
+    stage's output, which is empty when model preparation is skipped. **All
+    three are required once the quantization stage executes**, and strategies
+    should validate their presence.
 
     Flows needing no quantization (FP16, GPU backends) skip the stage entirely
     via ``GenAIPipeline.from_proxy(proxy, skip_stages={STAGE_QUANTIZATION})``
@@ -36,8 +37,17 @@ class QuantizationInputConfig:
         soc_model: The target SoC (e.g., QcomChipset.SM8750). Required.
         backend_type: QNN backend type (HTP, GPU, LPAI, etc.). Required.
         model_module: The nn.Module to quantize. Required when the stage runs.
+        example_inputs: Positional example inputs for ``torch.export``. Required
+            when the stage runs. Sourced from the **model** via
+            ``ModelLoaderAdapter.get_example_inputs``, never from
+            ``calibration_data``: this tuple defines the exported graph's
+            positional signature, supplies the zero-initialized KV caches a
+            dataset sample does not carry, and fixes the AR length because HTP
+            has no dynamic shapes.
         calibration_data: Calibration samples. Required when the stage runs. Any
-            Iterable[Tuple[Tensor, ...]], including a DataLoader.
+            Iterable[Tuple[Tensor, ...]], including a DataLoader. Consumed only
+            by ``calibrate()`` -- it is never indexed or peeked at, so a
+            single-use generator stays intact.
         training_data: Training dataset for quantization-aware training (QAT),
             typically (features, labels) pairs. Mirrors ``qat_training_data`` in
             ``build_executorch_binary``. ``None`` selects PTQ.
@@ -48,6 +58,7 @@ class QuantizationInputConfig:
     soc_model: "QcomChipset"
     backend_type: "QnnExecuTorchBackendType"
     model_module: Optional["nn.Module"] = None
+    example_inputs: Optional[Tuple[Any, ...]] = None
     calibration_data: Optional[Iterable[Any]] = None
     training_data: Optional[Iterable[Any]] = None
     quant_recipe: Any = None

backends/qualcomm/genai_pipeline/genai_pipeline.py

@@ -213,6 +213,8 @@ def _run_quantization(
                 soc_model=context.soc_model,
                 backend_type=self._engine_proxy.backend_type,
                 model_module=model_prep_output.model_module,
+                # Export inputs come from the model, not from calibration_data.
+                example_inputs=model_prep_output.example_inputs,
                 calibration_data=model_prep_output.calibration_data,
             )
             output = self._quantization_stage.invoke(context, input_config)

backends/qualcomm/genai_pipeline/strategies/inference/default_device_runner_adapter.py

@@ -77,13 +77,18 @@ def execute(
     ) -> InferenceResult:
         """Execute the model on device via ADB.
 
+        Note: This is a two-step protocol. ``execute()`` runs the model and
+        returns performance metrics. Call ``pull_results()`` afterward to
+        retrieve the actual output data files from the device.
+
         Args:
             inference_options: Engine-specific options. Supported keys:
                 - ``method_index``: Index of the method to execute (default 0).
                 - ``iteration``: Number of inference iterations (default 1).
 
         Returns:
-            InferenceResult with output data and performance metrics.
+            InferenceResult with performance metrics. ``output_data`` is None
+            until ``pull_results()`` is called separately.
         """
         if self._adb is None:
             raise RuntimeError(

backends/qualcomm/genai_pipeline/strategies/model_preparation/__init__.py

@@ -4,6 +4,9 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.executorch_model_preparation_strategy import (
+    ExecuTorchModelPreparationStrategy,
+)
 from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.model_loader_adapter import (
     ModelLoaderAdapter,
 )
@@ -12,6 +15,7 @@
 )
 
 __all__ = [
+    "ExecuTorchModelPreparationStrategy",
     "ModelLoaderAdapter",
     "ModelPreparationStrategy",
 ]

backends/qualcomm/genai_pipeline/strategies/model_preparation/default_model_loader_adapter.py

@@ -8,7 +8,7 @@
 
 import logging
 from pathlib import Path
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Tuple
 
 logger = logging.getLogger(__name__)
 
@@ -30,6 +30,17 @@ class DefaultModelLoaderAdapter:
         ``CalibrationDataAdapter``.
     """
 
+    #: Batch size and sequence length of the generated example inputs. HTP has no
+    #: dynamic shapes, so these dimensions are baked into the exported graph.
+    DEFAULT_BATCH_SIZE = 1
+    DEFAULT_AR_LEN = 1
+
+    #: Preferred runtime tokenizer file names, in priority order.
+    #: ``pytorch_tokenizers.get_tokenizer`` dispatches on the file extension
+    #: (``.json`` -> ``HuggingFaceTokenizer``, otherwise Llama2c/Tiktoken), so the
+    #: file we hand back selects the runtime tokenizer implementation.
+    RUNTIME_TOKENIZER_NAMES = ("tokenizer.json", "tokenizer.model")
+
     def load_model(
         self,
         model_name: str,
@@ -94,6 +105,51 @@ def load_tokenizer(
         logger.info("Tokenizer loaded successfully")
         return tokenizer
 
+    def get_example_inputs(
+        self,
+        model: Any,
+        extra_options: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[Any, ...]:
+        """Build example inputs for ``torch.export`` from the model itself.
+
+        Prefers the model's own ``get_example_inputs()`` when it exposes one, so
+        models that already describe their export signature (the LLM wrappers
+        build a flat ``(tokens, attn_mask, pos_ids, *k_caches, *v_caches)``
+        tuple) stay authoritative. Otherwise a minimal ``(input_ids,)`` is
+        synthesized, which is the correct signature for a plain HuggingFace
+        causal LM without an external KV cache.
+
+        Args:
+            model: The module returned by :meth:`load_model`.
+            extra_options: Additional options. Supported keys:
+                - ``batch_size``: Batch dimension (default:
+                  ``DEFAULT_BATCH_SIZE``).
+                - ``ar_len``: Sequence length / autoregressive window
+                  (default: ``DEFAULT_AR_LEN``).
+
+        Returns:
+            A flat tuple positionally matching ``model.forward``.
+        """
+        import torch
+
+        extra_options = extra_options or {}
+
+        model_provided = getattr(model, "get_example_inputs", None)
+        if callable(model_provided):
+            logger.info("Using example inputs provided by the model")
+            return tuple(model_provided())
+
+        batch_size = extra_options.get("batch_size", self.DEFAULT_BATCH_SIZE)
+        ar_len = extra_options.get("ar_len", self.DEFAULT_AR_LEN)
+
+        logger.info(
+            "Synthesizing example inputs with batch_size=%d, ar_len=%d",
+            batch_size,
+            ar_len,
+        )
+        # int64 token ids: the embedding lookup indexes with them.
+        return (torch.zeros((batch_size, ar_len), dtype=torch.int64),)
+
     def export_tokenizer(
         self,
         tokenizer: Any,
@@ -102,11 +158,19 @@ def export_tokenizer(
     ) -> Path:
         """Export tokenizer to disk and return the runtime tokenizer file.
 
-        ``save_pretrained`` writes several files and returns the tuple of paths
-        it wrote, with the tokenizer file last. Both ``llm::load_tokenizer`` and
-        ``pytorch_tokenizers.get_tokenizer`` expect that **single file**, not the
-        containing directory, so we return it -- mirroring the existing
-        ``TokenizerWrapper._from_hf`` flow.
+        ``save_pretrained`` writes several files and returns the tuple of paths it
+        wrote. Both ``llm::load_tokenizer`` and ``pytorch_tokenizers.get_tokenizer``
+        expect a **single file**, not the containing directory, so one artifact has
+        to be singled out.
+
+        The file is chosen **by name** -- ``tokenizer.json`` first, then
+        ``tokenizer.model`` -- rather than by position in the returned tuple.
+        ``get_tokenizer`` dispatches on the extension, so picking the wrong
+        artifact silently constructs the wrong tokenizer class instead of raising,
+        and ``save_pretrained``'s ordering is an implementation detail that varies
+        with the tokenizer (fast vs slow, whether ``added_tokens.json`` is
+        written). ``artifacts[-1]`` remains a last-resort fallback for tokenizers
+        that emit neither name, mirroring ``TokenizerWrapper._from_hf``.
 
         Args:
             tokenizer: The tokenizer instance to export.
@@ -128,6 +192,31 @@ def export_tokenizer(
                 f"save_pretrained() reported no tokenizer artifacts in {output_dir}."
             )
 
-        runtime_tokenizer_path = Path(artifacts[-1])
+        runtime_tokenizer_path = self._select_runtime_tokenizer(artifacts)
         logger.info("Tokenizer exported to %s", runtime_tokenizer_path)
         return runtime_tokenizer_path
+
+    @classmethod
+    def _select_runtime_tokenizer(cls, artifacts: Any) -> Path:
+        """Pick the runtime tokenizer file out of ``save_pretrained``'s artifacts.
+
+        Args:
+            artifacts: The paths reported by ``save_pretrained``.
+
+        Returns:
+            The first artifact matching :attr:`RUNTIME_TOKENIZER_NAMES`, falling
+            back to the last artifact when none matches.
+        """
+        paths = [Path(artifact) for artifact in artifacts]
+
+        for name in cls.RUNTIME_TOKENIZER_NAMES:
+            for path in paths:
+                if path.name == name:
+                    return path
+
+        logger.warning(
+            "None of %s found among tokenizer artifacts; falling back to %s.",
+            ", ".join(cls.RUNTIME_TOKENIZER_NAMES),
+            paths[-1],
+        )
+        return paths[-1]