From e3bdc682ec5db6728993b6622ac50cfb8dab4de1 Mon Sep 17 00:00:00 2001 From: Arik Horodniceanu Date: Tue, 28 Apr 2026 13:26:56 -0700 Subject: [PATCH] Qualcomm AI Engine Direct - [GenAI Pipeline] PR3: Pipeline Orchestrator --- backends/qualcomm/genai_pipeline/__init__.py | 1 + .../qualcomm/genai_pipeline/genai_pipeline.py | 264 ++++++++++++++++ .../executorch_model_preparation_strategy.py | 43 +++ .../tests/test_genai_pipeline.py | 290 ++++++++++++++++++ 4 files changed, 598 insertions(+) create mode 100644 backends/qualcomm/genai_pipeline/genai_pipeline.py create mode 100644 backends/qualcomm/genai_pipeline/strategies/model_preparation/executorch_model_preparation_strategy.py create mode 100644 backends/qualcomm/genai_pipeline/tests/test_genai_pipeline.py diff --git a/backends/qualcomm/genai_pipeline/__init__.py b/backends/qualcomm/genai_pipeline/__init__.py index 9ff124fa7af..57d2eb09230 100644 --- a/backends/qualcomm/genai_pipeline/__init__.py +++ b/backends/qualcomm/genai_pipeline/__init__.py @@ -36,6 +36,7 @@ "EngineNotAvailableError", "EngineProxy", "EngineType", + "GenAIPipeline", "InferenceInputConfig", "InferenceOutputConfig", "ModelPreparationInputConfig", diff --git a/backends/qualcomm/genai_pipeline/genai_pipeline.py b/backends/qualcomm/genai_pipeline/genai_pipeline.py new file mode 100644 index 00000000000..8092bcb13b5 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/genai_pipeline.py @@ -0,0 +1,264 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import logging +import time +from pathlib import Path +from typing import Dict, Optional, Set, Tuple, Type + +from executorch.backends.qualcomm.genai_pipeline.configs.compilation_input_config import ( + CompilationInputConfig, +) +from executorch.backends.qualcomm.genai_pipeline.configs.compilation_output_config import ( + CompilationOutputConfig, +) +from executorch.backends.qualcomm.genai_pipeline.configs.inference_input_config import ( + InferenceInputConfig, +) +from executorch.backends.qualcomm.genai_pipeline.configs.inference_output_config import ( + InferenceOutputConfig, +) +from executorch.backends.qualcomm.genai_pipeline.configs.model_preparation_input_config import ( + ModelPreparationInputConfig, +) +from executorch.backends.qualcomm.genai_pipeline.configs.model_preparation_output_config import ( + ModelPreparationOutputConfig, +) +from executorch.backends.qualcomm.genai_pipeline.configs.quantization_input_config import ( + QuantizationInputConfig, +) +from executorch.backends.qualcomm.genai_pipeline.configs.quantization_output_config import ( + QuantizationOutputConfig, +) +from executorch.backends.qualcomm.genai_pipeline.engine_proxy import EngineProxy +from executorch.backends.qualcomm.genai_pipeline.pipeline_context import PipelineContext +from executorch.backends.qualcomm.genai_pipeline.pipeline_types import ( + EngineType, + STAGE_COMPILATION, + STAGE_INFERENCE, + STAGE_MODEL_PREPARATION, + STAGE_QUANTIZATION, +) +from executorch.backends.qualcomm.genai_pipeline.stages.compilation_stage import ( + CompilationStage, +) +from executorch.backends.qualcomm.genai_pipeline.stages.inference_stage import ( + InferenceStage, +) +from executorch.backends.qualcomm.genai_pipeline.stages.model_preparation_stage import ( + ModelPreparationStage, +) +from executorch.backends.qualcomm.genai_pipeline.stages.quantization_stage import ( + QuantizationStage, +) +from executorch.backends.qualcomm.genai_pipeline.strategies.compilation.executorch_compilation_strategy import ( + ExecuTorchCompilationStrategy, +) +from executorch.backends.qualcomm.genai_pipeline.strategies.inference.executorch_inference_strategy import ( + ExecuTorchInferenceStrategy, +) +from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.executorch_model_preparation_strategy import ( + ExecuTorchModelPreparationStrategy, +) +from executorch.backends.qualcomm.genai_pipeline.strategies.quantization.executorch_quantization_strategy import ( + ExecuTorchQuantizationStrategy, +) + +logger = logging.getLogger("genai_pipeline") + +# Maps (stage_name, engine_type) → (StageClass, StrategyClass) +_STRATEGY_REGISTRY: Dict[Tuple[str, EngineType], Tuple[Type, Type]] = { + (STAGE_MODEL_PREPARATION, EngineType.EXECUTORCH): ( + ModelPreparationStage, + ExecuTorchModelPreparationStrategy, + ), + (STAGE_QUANTIZATION, EngineType.EXECUTORCH): ( + QuantizationStage, + ExecuTorchQuantizationStrategy, + ), + (STAGE_COMPILATION, EngineType.EXECUTORCH): ( + CompilationStage, + ExecuTorchCompilationStrategy, + ), + (STAGE_INFERENCE, EngineType.EXECUTORCH): ( + InferenceStage, + ExecuTorchInferenceStrategy, + ), +} + + +class GenAIPipeline: + """Main pipeline orchestrator for GenAI model workflows. + + Assembles stages from EngineProxy, wires data flow between + InputConfig → Stage → OutputConfig, and executes sequentially: + model_preparation → quantization → compilation → inference. + """ + + def __init__( + self, + model_preparation_stage: Optional[ModelPreparationStage], + quantization_stage: Optional[QuantizationStage], + compilation_stage: Optional[CompilationStage], + inference_stage: Optional[InferenceStage], + engine_proxy: EngineProxy, + ) -> None: + self._model_preparation_stage = model_preparation_stage + self._quantization_stage = quantization_stage + self._compilation_stage = compilation_stage + self._inference_stage = inference_stage + self._engine_proxy = engine_proxy + + @classmethod + def from_proxy( + cls, + engine_proxy: EngineProxy, + skip_stages: Optional[Set[str]] = None, + ) -> "GenAIPipeline": + """Create a pipeline from an EngineProxy configuration. + + Args: + engine_proxy: The EngineProxy with per-stage engine assignments. + skip_stages: Optional set of stage names to skip. + + Returns: + A fully configured GenAIPipeline. + """ + skip = skip_stages or set() + + model_prep_stage = cls._resolve_stage( + engine_proxy, STAGE_MODEL_PREPARATION, skip + ) + quant_stage = cls._resolve_stage(engine_proxy, STAGE_QUANTIZATION, skip) + compile_stage = cls._resolve_stage(engine_proxy, STAGE_COMPILATION, skip) + infer_stage = cls._resolve_stage(engine_proxy, STAGE_INFERENCE, skip) + + return cls( + model_preparation_stage=model_prep_stage, + quantization_stage=quant_stage, + compilation_stage=compile_stage, + inference_stage=infer_stage, + engine_proxy=engine_proxy, + ) + + @staticmethod + def _resolve_stage( + engine_proxy: EngineProxy, + stage_name: str, + skip: set, + ): + if stage_name in skip: + return None + + engine = engine_proxy.get_engine(stage_name) + key = (stage_name, engine) + + if key not in _STRATEGY_REGISTRY: + raise ValueError(f"No {stage_name} strategy for engine: {engine}") + + stage_cls, strategy_cls = _STRATEGY_REGISTRY[key] + return stage_cls(strategy_cls()) + + def invoke(self, context: PipelineContext) -> InferenceOutputConfig: + """Execute the full pipeline sequentially. + + Args: + context: The pipeline context with user inputs. + + Returns: + InferenceOutputConfig with inference results and metrics. + """ + logger.info( + "[GenAIPipeline] Pipeline started for model '%s'", context.model_name + ) + + model_prep_output = self._run_model_preparation(context) + quant_output = self._run_quantization(context, model_prep_output) + compile_output = self._run_compilation(context, quant_output) + result = self._run_inference(context, model_prep_output, compile_output) + + logger.info("[GenAIPipeline] Pipeline completed") + return result + + def _run_model_preparation( + self, context: PipelineContext + ) -> ModelPreparationOutputConfig: + if self._model_preparation_stage is not None: + logger.info("[GenAIPipeline] ModelPreparationStage started") + start = time.monotonic() + input_config = ModelPreparationInputConfig( + model_name=context.model_name, + soc_model=context.soc_model, + ) + output = self._model_preparation_stage.invoke(context, input_config) + elapsed = time.monotonic() - start + logger.info( + "[GenAIPipeline] ModelPreparationStage completed in %.1fs", elapsed + ) + return output + return ModelPreparationOutputConfig() + + def _run_quantization( + self, + context: PipelineContext, + model_prep_output: ModelPreparationOutputConfig, + ) -> QuantizationOutputConfig: + if self._quantization_stage is not None: + logger.info("[GenAIPipeline] QuantizationStage started") + start = time.monotonic() + input_config = QuantizationInputConfig( + soc_model=context.soc_model, + backend_type=self._engine_proxy.backend_type, + model_module=model_prep_output.model_module, + calibration_data=model_prep_output.calibration_data, + ) + output = self._quantization_stage.invoke(context, input_config) + elapsed = time.monotonic() - start + logger.info("[GenAIPipeline] QuantizationStage completed in %.1fs", elapsed) + return output + return QuantizationOutputConfig() + + def _run_compilation( + self, + context: PipelineContext, + quant_output: QuantizationOutputConfig, + ) -> CompilationOutputConfig: + if self._compilation_stage is not None: + logger.info("[GenAIPipeline] CompilationStage started") + start = time.monotonic() + input_config = CompilationInputConfig( + soc_model=context.soc_model, + backend_type=self._engine_proxy.backend_type, + model=quant_output.quantized_model, + artifact_dir=Path(context.artifact_dir), + ) + output = self._compilation_stage.invoke(context, input_config) + elapsed = time.monotonic() - start + logger.info("[GenAIPipeline] CompilationStage completed in %.1fs", elapsed) + return output + return CompilationOutputConfig() + + def _run_inference( + self, + context: PipelineContext, + model_prep_output: ModelPreparationOutputConfig, + compile_output: CompilationOutputConfig, + ) -> InferenceOutputConfig: + if self._inference_stage is not None: + logger.info("[GenAIPipeline] InferenceStage started") + start = time.monotonic() + input_config = InferenceInputConfig( + artifact_paths=compile_output.artifact_paths, + tokenizer=model_prep_output.tokenizer, + runtime_tokenizer_path=model_prep_output.runtime_tokenizer_path, + prompt=context.prompt, + soc_model=context.soc_model, + ) + output = self._inference_stage.invoke(context, input_config) + elapsed = time.monotonic() - start + logger.info("[GenAIPipeline] InferenceStage completed in %.1fs", elapsed) + return output + return InferenceOutputConfig() diff --git a/backends/qualcomm/genai_pipeline/strategies/model_preparation/executorch_model_preparation_strategy.py b/backends/qualcomm/genai_pipeline/strategies/model_preparation/executorch_model_preparation_strategy.py new file mode 100644 index 00000000000..8322f7a68b6 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/strategies/model_preparation/executorch_model_preparation_strategy.py @@ -0,0 +1,43 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# 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.configs.model_preparation_input_config import ( + ModelPreparationInputConfig, +) +from executorch.backends.qualcomm.genai_pipeline.configs.model_preparation_output_config import ( + ModelPreparationOutputConfig, +) +from executorch.backends.qualcomm.genai_pipeline.pipeline_context import PipelineContext +from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.model_preparation_strategy import ( + ModelPreparationStrategy, +) + + +class ExecuTorchModelPreparationStrategy(ModelPreparationStrategy): + """ExecuTorch-based model preparation using HuggingFace transformers. + + Loads model weights, tokenizer, and generates calibration data + for downstream quantization and compilation stages. + """ + + def invoke( + self, + context: PipelineContext, + input_config: ModelPreparationInputConfig, + ) -> ModelPreparationOutputConfig: + """Prepare the model, tokenizer, and calibration data. + + Args: + context: The pipeline context. + input_config: The model preparation input configuration. + + Returns: + ModelPreparationOutputConfig with model, tokenizer, and calibration data. + """ + raise NotImplementedError( + "ExecuTorchModelPreparationStrategy.invoke() is a stub. " + "Implementation will be added in a subsequent PR." + ) diff --git a/backends/qualcomm/genai_pipeline/tests/test_genai_pipeline.py b/backends/qualcomm/genai_pipeline/tests/test_genai_pipeline.py new file mode 100644 index 00000000000..1083950035b --- /dev/null +++ b/backends/qualcomm/genai_pipeline/tests/test_genai_pipeline.py @@ -0,0 +1,290 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest +from unittest.mock import MagicMock + +from executorch.backends.qualcomm.genai_pipeline.configs.compilation_output_config import ( + CompilationOutputConfig, +) +from executorch.backends.qualcomm.genai_pipeline.configs.inference_output_config import ( + InferenceOutputConfig, +) +from executorch.backends.qualcomm.genai_pipeline.configs.quantization_output_config import ( + QuantizationOutputConfig, +) +from executorch.backends.qualcomm.genai_pipeline.engine_proxy import EngineProxy +from executorch.backends.qualcomm.genai_pipeline.genai_pipeline import GenAIPipeline +from executorch.backends.qualcomm.genai_pipeline.pipeline_types import ( + EngineType, + STAGE_COMPILATION, + STAGE_INFERENCE, + STAGE_MODEL_PREPARATION, + STAGE_QUANTIZATION, +) +from executorch.backends.qualcomm.genai_pipeline.stages.compilation_stage import ( + CompilationStage, +) +from executorch.backends.qualcomm.genai_pipeline.stages.inference_stage import ( + InferenceStage, +) +from executorch.backends.qualcomm.genai_pipeline.stages.quantization_stage import ( + QuantizationStage, +) +from executorch.backends.qualcomm.genai_pipeline.strategies.compilation.compilation_strategy import ( + CompilationStrategy, +) +from executorch.backends.qualcomm.genai_pipeline.strategies.inference.inference_strategy import ( + InferenceStrategy, +) +from executorch.backends.qualcomm.genai_pipeline.strategies.quantization.quantization_strategy import ( + QuantizationStrategy, +) +from executorch.backends.qualcomm.genai_pipeline.tests.test_utils import ( + make_test_context, + TEST_PTE_PATH, +) + +TEST_MOCK_GENERATED_TEXT = "Mock generated text" +TEST_MOCK_TOKENS_PER_SEC = 42.0 +TEST_MOCK_BACKEND_TYPE = MagicMock(name="kHtpBackend") + + +class _MockQuantizationStrategy(QuantizationStrategy): + def invoke(self, context, input_config): + return QuantizationOutputConfig(quantized_model="mock_quantized_model") + + +class _MockCompilationStrategy(CompilationStrategy): + def invoke(self, context, input_config): + return CompilationOutputConfig(artifact_paths=[TEST_PTE_PATH]) + + +class _MockInferenceStrategy(InferenceStrategy): + def invoke(self, context, input_config): + return InferenceOutputConfig( + inference_results=[TEST_MOCK_GENERATED_TEXT], + performance_metrics={"tokens_per_sec": TEST_MOCK_TOKENS_PER_SEC}, + ) + + +class TestGenAIPipelineFromProxy(unittest.TestCase): + + def test_full_executorch_creates_all_stages(self): + proxy = EngineProxy( + { + STAGE_MODEL_PREPARATION: EngineType.EXECUTORCH, + STAGE_QUANTIZATION: EngineType.EXECUTORCH, + STAGE_COMPILATION: EngineType.EXECUTORCH, + STAGE_INFERENCE: EngineType.EXECUTORCH, + }, + backend_type=TEST_MOCK_BACKEND_TYPE, + ) + pipeline = GenAIPipeline.from_proxy(proxy) + self.assertIsNotNone(pipeline._model_preparation_stage) + self.assertIsNotNone(pipeline._quantization_stage) + self.assertIsNotNone(pipeline._compilation_stage) + self.assertIsNotNone(pipeline._inference_stage) + + def test_skip_stages(self): + proxy = EngineProxy( + {STAGE_INFERENCE: EngineType.EXECUTORCH}, + backend_type=TEST_MOCK_BACKEND_TYPE, + ) + pipeline = GenAIPipeline.from_proxy( + proxy, + skip_stages={ + STAGE_MODEL_PREPARATION, + STAGE_QUANTIZATION, + STAGE_COMPILATION, + STAGE_INFERENCE, + }, + ) + self.assertIsNone(pipeline._model_preparation_stage) + self.assertIsNone(pipeline._quantization_stage) + self.assertIsNone(pipeline._compilation_stage) + self.assertIsNone(pipeline._inference_stage) + + def test_default_engines_all_executorch(self): + proxy = EngineProxy({}, backend_type=TEST_MOCK_BACKEND_TYPE) + pipeline = GenAIPipeline.from_proxy(proxy) + self.assertIsNotNone(pipeline._model_preparation_stage) + self.assertIsNotNone(pipeline._quantization_stage) + self.assertIsNotNone(pipeline._compilation_stage) + self.assertIsNotNone(pipeline._inference_stage) + + def test_from_proxy_unsupported_model_preparation_engine_raises(self): + proxy = MagicMock(spec=EngineProxy) + proxy.get_engine.return_value = MagicMock() + with self.assertRaises(ValueError) as cm: + GenAIPipeline.from_proxy(proxy) + self.assertIn("No model_preparation strategy", str(cm.exception)) + + def test_from_proxy_unsupported_quantization_engine_raises(self): + proxy = MagicMock(spec=EngineProxy) + proxy.get_engine.side_effect = lambda stage: ( + EngineType.EXECUTORCH if stage == STAGE_MODEL_PREPARATION else MagicMock() + ) + with self.assertRaises(ValueError) as cm: + GenAIPipeline.from_proxy(proxy) + self.assertIn("No quantization strategy", str(cm.exception)) + + def test_from_proxy_unsupported_compilation_engine_raises(self): + proxy = MagicMock(spec=EngineProxy) + proxy.get_engine.side_effect = lambda stage: ( + EngineType.EXECUTORCH + if stage in (STAGE_MODEL_PREPARATION, STAGE_QUANTIZATION) + else MagicMock() + ) + with self.assertRaises(ValueError) as cm: + GenAIPipeline.from_proxy(proxy) + self.assertIn("No compilation strategy", str(cm.exception)) + + def test_from_proxy_unsupported_inference_engine_raises(self): + proxy = MagicMock(spec=EngineProxy) + proxy.get_engine.side_effect = lambda stage: ( + EngineType.EXECUTORCH + if stage in (STAGE_MODEL_PREPARATION, STAGE_QUANTIZATION, STAGE_COMPILATION) + else MagicMock() + ) + with self.assertRaises(ValueError) as cm: + GenAIPipeline.from_proxy(proxy) + self.assertIn("No inference strategy", str(cm.exception)) + + +class TestGenAIPipelineInvoke(unittest.TestCase): + + def test_invoke_with_mock_strategies(self): + proxy = EngineProxy( + { + STAGE_QUANTIZATION: EngineType.EXECUTORCH, + STAGE_COMPILATION: EngineType.EXECUTORCH, + STAGE_INFERENCE: EngineType.EXECUTORCH, + }, + backend_type=TEST_MOCK_BACKEND_TYPE, + ) + pipeline = GenAIPipeline( + model_preparation_stage=None, + quantization_stage=QuantizationStage(_MockQuantizationStrategy()), + compilation_stage=CompilationStage(_MockCompilationStrategy()), + inference_stage=InferenceStage(_MockInferenceStrategy()), + engine_proxy=proxy, + ) + result = pipeline.invoke(make_test_context()) + + self.assertIsInstance(result, InferenceOutputConfig) + self.assertEqual(result.inference_results, [TEST_MOCK_GENERATED_TEXT]) + self.assertEqual( + result.performance_metrics["tokens_per_sec"], TEST_MOCK_TOKENS_PER_SEC + ) + + def test_invoke_compile_only(self): + proxy = EngineProxy( + { + STAGE_QUANTIZATION: EngineType.EXECUTORCH, + STAGE_COMPILATION: EngineType.EXECUTORCH, + }, + backend_type=TEST_MOCK_BACKEND_TYPE, + ) + pipeline = GenAIPipeline( + model_preparation_stage=None, + quantization_stage=QuantizationStage(_MockQuantizationStrategy()), + compilation_stage=CompilationStage(_MockCompilationStrategy()), + inference_stage=None, + engine_proxy=proxy, + ) + result = pipeline.invoke(make_test_context()) + + self.assertIsInstance(result, InferenceOutputConfig) + + def test_invoke_no_stages(self): + proxy = EngineProxy({}, backend_type=TEST_MOCK_BACKEND_TYPE) + pipeline = GenAIPipeline( + model_preparation_stage=None, + quantization_stage=None, + compilation_stage=None, + inference_stage=None, + engine_proxy=proxy, + ) + result = pipeline.invoke(make_test_context()) + + self.assertIsInstance(result, InferenceOutputConfig) + + def test_quantization_receives_soc_model(self): + mock_quant = MagicMock(spec=QuantizationStrategy) + mock_quant.invoke.return_value = QuantizationOutputConfig( + quantized_model="quantized" + ) + + test_soc = "SM8650" + proxy = EngineProxy( + {STAGE_QUANTIZATION: EngineType.EXECUTORCH}, + backend_type=TEST_MOCK_BACKEND_TYPE, + ) + pipeline = GenAIPipeline( + model_preparation_stage=None, + quantization_stage=QuantizationStage(mock_quant), + compilation_stage=None, + inference_stage=None, + engine_proxy=proxy, + ) + pipeline.invoke(make_test_context(soc_model=test_soc)) + + args, _ = mock_quant.invoke.call_args + input_config = args[1] + self.assertEqual(input_config.soc_model, test_soc) + + def test_compilation_receives_backend_type(self): + mock_compile = MagicMock(spec=CompilationStrategy) + mock_compile.invoke.return_value = CompilationOutputConfig( + artifact_paths=[TEST_PTE_PATH] + ) + + mock_backend_type = MagicMock() + proxy = EngineProxy( + {STAGE_COMPILATION: EngineType.EXECUTORCH}, + backend_type=mock_backend_type, + ) + pipeline = GenAIPipeline( + model_preparation_stage=None, + quantization_stage=None, + compilation_stage=CompilationStage(mock_compile), + inference_stage=None, + engine_proxy=proxy, + ) + pipeline.invoke(make_test_context()) + + call_args = mock_compile.invoke.call_args + input_config = call_args[0][1] + self.assertEqual(input_config.backend_type, mock_backend_type) + + def test_inference_receives_prompt(self): + mock_infer = MagicMock(spec=InferenceStrategy) + mock_infer.invoke.return_value = InferenceOutputConfig( + inference_results=["output"] + ) + + test_prompt = ["What is AI?"] + proxy = EngineProxy( + {STAGE_INFERENCE: EngineType.EXECUTORCH}, + backend_type=TEST_MOCK_BACKEND_TYPE, + ) + pipeline = GenAIPipeline( + model_preparation_stage=None, + quantization_stage=None, + compilation_stage=None, + inference_stage=InferenceStage(mock_infer), + engine_proxy=proxy, + ) + pipeline.invoke(make_test_context(prompt=test_prompt)) + + call_args = mock_infer.invoke.call_args + input_config = call_args[0][1] + self.assertEqual(input_config.prompt, test_prompt) + + +if __name__ == "__main__": + unittest.main()