Qualcomm AI Engine Direct - [GenAI Pipeline Phase 2] PRB1 - Compilation foundation

[qti-horodnic]

Summary

First PR of Phase 2 and the start of Stream B (compilation). Phase 1 added the skeleton, interfaces, stage wrappers, typed configs with stubbed logic. This PR adds the compilation groundwork the rest of Stream B builds on and makes DefaultCompilerAdapter actually compile a graph to a .pte.

Nothing under examples/ changes. The legacy llama.py flow is untouched and stays CI's reference; the two coexist until Phase 3, whose deletion of the legacy path is gated on a parity test added in PR-A2.

What's included

Two separate name spaces: artifact_keys.py, graph_names.py

An artifact key names a .pte file; a graph name names a method inside one. A hybrid decoder produces one .pte (text_decoder) holding two methods (kv_forward, prefill_forward), so these are deliberately distinct key spaces.

• Artifact keys are the six the on-device runner already uses for pte_paths, so a compiled bundle reaches inference without translation. Plus DECODE_QDQ_FILENAME, a filename (a .pt2 the SQNR eval reads, never seen by the runner) and therefore deliberately absent from ALL_ARTIFACT_KEYS.
• DECODER_GRAPH_NAMES is decode first, and the order matters: kv mode builds no prefill graph, and decode is the authoritative source of the artifact's constant methods.
• Both sets duplicate strings from examples/.../decoder_constants.py rather than importing them, so this package does not depend on the example scripts. Agreement tests keep them in step; the legacy copies go away with the legacy flow.

graph_bundle.py: GraphBundle, the frozen per-graph unit quantization hands to compilation (module, inputs, meta, quant_io_dtypes, modality_inputs, executorch_config)

• One decoder is exported several times from the same weights, so the stages after model preparation work on a set of graphs. Bundling them means the pipeline threads one {graph_name: GraphBundle} map rather than several parallel dicts that can drift apart.
• frozen=True stops in-place edits across a stage boundary. Immutability stops at the field boundary though, convert_pt2e mutates meta in place and the docstring says so, so the guarantee is not overread.
• quant_io_dtypes is the one thing compilation cannot derive: the graph-boundary dtypes follow from the quantization recipe's bit widths. None means quantization was skipped. __post_init__ requires both keys or neither, because the tagger indexes both unconditionally so a partial mapping is a KeyError at lowering, not a partially-quantized boundary.
• CompilationInputConfig.graphs is the consumer side, added here; PR-A1 adds the producer side, PR-B2 is the first stage to read it. Defined here so both streams import one definition.

compilation/compile_spec_builder.py: QnnCompileSpecBuilder, resolve_soc_model, resolve_backend_type

• Pairs generate_htp_compiler_spec / generate_gpu_compiler_spec with generate_qnn_executorch_compiler_spec behind one build(); llama.py repeats that branch three times today.
• resolve_soc_model converts str → QcomChipset. PipelineContext carries the SoC as a string so nothing in the pipeline imports the QNN schema; lowering needs the enum. Per earlier review this belongs in the adapter layer.
• resolve_backend_type uses a lookup map built from the enum itself, since QnnExecuTorchBackendType.__str__ already yields the CLI name so no hardcoded table, and a bad name raises ValueError listing the valid ones.
• enable_x86_64 is constructor state, not a per-call flag: the emulator supports neither weight sharing nor shared buffers, which is exactly why llama.py writes not args.enable_x86_64 at all three call sites. Hardcoding either on would break the x86 CI run without failing anything locally.
• Trap worth knowing: the builder's device default for shared_buffer is on, while ControlArgs.shared_buffer defaults off to mirror llama.py's parser. A caller driving the builder from a ControlArgs must pass it explicitly. Documented on build() and pinned by two tests.

control_args.py: typed bridge from PipelineContext to the existing LLM components

• Those components read configuration off an argparse.Namespace (args.max_seq_len, args.enable_x86_64, …), so this is a dataclass carrying the same 65 attributes, buildable from a PipelineContext.
• It subclasses argparse.Namespace, and that is load-bearing: QnnConfig.load_config dispatches on isinstance(config, argparse.Namespace), so a plain dataclass would break the device path.
• build_parser() derives an argparse parser from the same fields, so a CLI entry point stops maintaining a second default table.
• Defaults are asserted field-for-field against the real llama.py parser by a test, since this is a second configuration surface and drift would be silent. Two fields are explicitly exempt (train_config, lr_config default to None rather than to YAML paths under examples/), with a test pinning that the exemption is deliberate.
• Documented as temporary: fields leave as their consumers move behind pipeline interfaces.

artifact_paths: List[Path] → Dict[str, Path] (CompilationResult, CompilationOutputConfig, InferenceInputConfig, DeviceRunnerAdapter, DefaultDeviceRunnerAdapter)

Keyed because the runner addresses artifacts individually and which ones exist varies by model, so position carries no reliable meaning. Absent artifacts are omitted, never mapped to None so the runner tests key membership to decide whether a model is multimodal. This is the frozen B→A2 interface.

DefaultCompilerAdapter implementation: replaces PR6's NotImplementedError

• Lowers the graph, converts with the same config the legacy path uses (MemoryPlanningPass(alloc_graph_input=False, alloc_graph_output=False) + BuildQuantIo(), since with a shared buffer the caller supplies graph I/O addresses from RPC memory), writes the .pte, and returns it under the artifact_key the caller names keyed by the caller because a 1:1 adapter cannot tell a vision encoder from a text decoder when both arrive as model.
• soc_model and backend_type are validated against compile_specs rather than trusted. Lowering takes its target only from the specs, so these parameters would otherwise be decorative: a caller could validate ops for one SoC and compile for another and see nothing until the artifact ran on device. The target is read back out of the specs with flatbuffer_to_option and a disagreement raises. Specs with no QNN entry skip the check.
• Two things stay deliberately outside the adapter, stated in its docstring: multi-graph grouping (fanning out into a multi-method .pte is the strategy's job, so the adapter stays a 1:1 wrapper) and spill-fill sizing (a property of the group, computed from the lowered program). Both land in PR-B2.

Two bugs found along the way (worth a look, not specific to this feature)

• DefaultDeviceRunnerAdapter.push_artifacts did [str(p) for p in artifact_paths]. Once that is a dict, iteration yields keys, so it would have pushed the strings "text_decoder", "tok_embedding" … instead of paths. Now .values(), matching EvalBase._get_adb. This is why the type change could not be a pure annotation edit.
• CompilationInputConfig imported CompileSpec from executorch.exir.backend.compile_spec, which does not exist (it is compile_spec_schema). Hidden at runtime by TYPE_CHECKING; mypy flags it.

Tests: 6 new files, 8 updated. Beyond the unit tests, DefaultCompilerAdapter was run against the real stack (a small module compiled through QNN to a non-empty .pte on the x86 path); that takes ~15s so the committed adapter tests mock lowering, while the builder tests keep two real-API cases (specs really are CompileSpecs; use_multi_contexts + online_prepare is rejected rather than masked).

PR Review Checklist

• All new classes follow single responsibility (one class per file) - Yes.
• All dependencies are injected via constructor with sensible defaults - Yes.
• All external calls are behind injectable interfaces - Yes (adapter pattern; lowering is imported lazily inside the adapter).
• Unit tests cover every public method - Yes.
• Docstrings on all public classes and methods - Yes.
• Type annotations on all function signatures - Yes.
• Logging follows the strategy in the LLD - Yes (info on entry/exit, debug per step, warning on degraded paths).
• No existing behaviour changed - Yes (no examples/ files touched; llama.py and the wrappers import unchanged).

Related PRs

Phase 2 flow. PR-A1 and PR-B1 are independent and can land in either order. PR-B2 depends on PR-B1. PR-A2 depends on all three, and closes Phase 2 by wiring the runner and adding the parity test that gates the Phase-3 deletion of the legacy flow.

PR-A1  ─────────────────────────┐
                                ├─► PR-A2
PR-B1 (this PR) ──► PR-B2 ──────┘

Phase 1 (merged):

• PR 1: Core data model, engine routing & exceptions: Qualcomm AI Engine Direct - [GenAI Pipeline] PR1: Core Data Model, Engine Routing & Exceptions #20409
• PR 2: Strategy interfaces & stage wrappers: Qualcomm AI Engine Direct - [GenAI Pipeline] PR2: Strategy Interfaces & Stage Wrappers #20795
• PR 3: Pipeline orchestrator: Qualcomm AI Engine Direct - [GenAI Pipeline] PR3: Pipeline Orchestrator #21149
• PR 4: Adapter interfaces & default implementations: Qualcomm AI Engine Direct - [GenAI Pipeline] PR4: Adapter interfaces, default implementations & dataset providers #21751
• PR 5: Model preparation & quantization strategies: Qualcomm AI Engine Direct - [GenAI Pipeline] PR5: Model preparation & quantization strategy implementations #21899
• PR 6: Compilation & inference strategy implementations: Qualcomm AI Engine Direct - [GenAI Pipeline] PR6: Compilation & inference strategy implementations #22284
• PR 7: Integration & E2E tests: cancelled, will be raised after Phase 2 is done.

Phase 2:

• PR A1: Model registry, model_lookup, CLI, LLMQuantizerAdapter, multi-graph quantization & encoding reconciliation. Consumes ControlArgs, GraphBundle and graph_names from this PR, but does not depend on it landing first; pending
• PR B1: Compilation foundation (this PR)
• PR B2: Multi-graph lowering, weight sharing, sharding & spill-fill. Depends on this PR; pending
• PR A2: Device-runner adapter, pipeline-runner wiring, README & the llama_stories_260k E2E parity test. Depends on PR-A1 and PR-B2; pending

Test plan

Run only tests added in this PR:

python -m pytest \
  backends/qualcomm/genai_pipeline/tests/test_artifact_keys.py \
  backends/qualcomm/genai_pipeline/tests/test_control_args.py \
  backends/qualcomm/genai_pipeline/tests/compilation/ \
  backends/qualcomm/genai_pipeline/tests/strategies/compilation/test_default_compiler_adapter.py \
  -v

Result:

72 passed, 76 subtests passed in 7.55s 

Run all genai_pipeline tests:

python -m pytest backends/qualcomm/genai_pipeline/tests/ -v

Result:

246 passed, 92 subtests passed in 8.08s

Run all tests with coverage:

python -m pytest backends/qualcomm/genai_pipeline/tests/ \
  --cov=backends/qualcomm/genai_pipeline \
  --cov-config=backends/qualcomm/.coveragerc \
  --cov-report=term-missing

Result:

Name                                                                                                     Stmts   Miss Branch BrPart  Cover   Missing
----------------------------------------------------------------------------------------------------------------------------------------------------
backends/qualcomm/genai_pipeline/artifact_keys.py                                                            9      0      0      0   100%
backends/qualcomm/genai_pipeline/compilation/compile_spec_builder.py                                        55      1     16      0    99%   181
backends/qualcomm/genai_pipeline/configs/compilation_input_config.py                                        13      0      0      0   100%
backends/qualcomm/genai_pipeline/configs/compilation_output_config.py                                        8      0      0      0   100%
backends/qualcomm/genai_pipeline/configs/inference_input_config.py                                          12      0      0      0   100%
backends/qualcomm/genai_pipeline/configs/inference_output_config.py                                          9      0      0      0   100%
backends/qualcomm/genai_pipeline/configs/model_preparation_input_config.py                                   7      0      0      0   100%
backends/qualcomm/genai_pipeline/configs/model_preparation_output_config.py                                 12      0      0      0   100%
backends/qualcomm/genai_pipeline/configs/quantization_input_config.py                                       13      0      0      0   100%
backends/qualcomm/genai_pipeline/configs/quantization_output_config.py                                       5      0      0      0   100%
backends/qualcomm/genai_pipeline/control_args.py                                                           154      1     24      1    99%   276
backends/qualcomm/genai_pipeline/datasets/calibration_data_adapter.py                                        5      0      0      0   100%
backends/qualcomm/genai_pipeline/datasets/default_calibration_data_adapter.py                               26      0      6      0   100%
backends/qualcomm/genai_pipeline/datasets/default_training_data_adapter.py                                  14      0      2      0   100%
backends/qualcomm/genai_pipeline/datasets/training_data_adapter.py                                           5      0      0      0   100%
backends/qualcomm/genai_pipeline/engine_proxy.py                                                            20      0      4      0   100%
backends/qualcomm/genai_pipeline/exceptions.py                                                              20      0      6      0   100%
backends/qualcomm/genai_pipeline/genai_pipeline.py                                                          99      7     12      1    93%   218-229
backends/qualcomm/genai_pipeline/graph_bundle.py                                                            18      0      4      0   100%
backends/qualcomm/genai_pipeline/graph_names.py                                                              8      0      0      0   100%
backends/qualcomm/genai_pipeline/pipeline_context.py                                                        52      0     14      0   100%
backends/qualcomm/genai_pipeline/pipeline_stage.py                                                           5      0      0      0   100%
backends/qualcomm/genai_pipeline/stages/compilation_stage.py                                                14      0      0      0   100%
backends/qualcomm/genai_pipeline/stages/inference_stage.py                                                  14      0      0      0   100%
backends/qualcomm/genai_pipeline/stages/model_preparation_stage.py                                          14      2      0      0    86%   30, 37
backends/qualcomm/genai_pipeline/stages/quantization_stage.py                                               14      0      0      0   100%
backends/qualcomm/genai_pipeline/strategies/compilation/compilation_strategy.py                              7      0      0      0   100%
backends/qualcomm/genai_pipeline/strategies/compilation/compiler_adapter.py                                 12      0      0      0   100%
backends/qualcomm/genai_pipeline/strategies/compilation/executorch_compilation_strategy.py                  45      0     10      0   100%
backends/qualcomm/genai_pipeline/strategies/inference/device_runner_adapter.py                              14      0      0      0   100%
backends/qualcomm/genai_pipeline/strategies/inference/executorch_inference_strategy.py                      43      0      6      0   100%
backends/qualcomm/genai_pipeline/strategies/inference/inference_strategy.py                                  7      0      0      0   100%
backends/qualcomm/genai_pipeline/strategies/model_preparation/executorch_model_preparation_strategy.py      68      0     14      0   100%
backends/qualcomm/genai_pipeline/strategies/model_preparation/model_loader_adapter.py                        9      0      0      0   100%
backends/qualcomm/genai_pipeline/strategies/model_preparation/model_preparation_strategy.py                  7      0      0      0   100%
backends/qualcomm/genai_pipeline/strategies/quantization/executorch_quantization_strategy.py                61      0     18      0   100%
backends/qualcomm/genai_pipeline/strategies/quantization/quantization_strategy.py                            7      0      0      0   100%
backends/qualcomm/genai_pipeline/strategies/quantization/quantizer_adapter.py                                9      0      0      0   100%
----------------------------------------------------------------------------------------------------------------------------------------------------
TOTAL                                                                                                      914     11    136      2    99%

Confirm the legacy flow is unaffected:

python -c "
from executorch.examples.qualcomm.oss_scripts.llama.llama import _build_parser, export_llama
from executorch.examples.qualcomm.oss_scripts.llama.wrappers import MultiModalManager, HybridAttentionSinkEvictor
print('LEGACY IMPORTS OK')
"

Result:

`LEGACY IMPORTS OK`

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22846

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 6748db8 with merge base e5acd53 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[qti-horodnic]

@pytorchbot label "release notes: qualcomm"

[psiddh]

@claude Review this pr , and share a high level and technical summary on how this improves underlying things and sets up foundation

[claude]

Claude finished @psiddh's task in 2m 2s —— View job

---

PR review: GenAI Pipeline Phase 2, PRB1 (compilation foundation)

• Gathered context (diff, CLAUDE.md, the pipeline and strategy call sites)
• Reviewed the new modules: artifact_keys, graph_names, graph_bundle, control_args, compile_spec_builder
• Reviewed DefaultCompilerAdapter and the List → Dict change for artifact_paths
• Posted the summary and review

---

High-level summary

Phase 1 built the pipeline's structure (stages, strategies, adapters, typed configs), but compilation was a stub: DefaultCompilerAdapter raised NotImplementedError. This PR is the first time the pipeline can turn a graph into a real .pte. It also sets up the shared types and interfaces the other Phase 2 PRs build on.

How it improves things:

1. Fixed names for artifacts and graphs. Before this PR, file and method names were loose strings copied from examples/. Now there are two separate name sets:

   • artifact_keys names the .pte files.
   • graph_names names the methods inside one .pte.

   One hybrid decoder .pte holds both kv_forward and prefill_forward, so the two need different names. Keeping them apart stops a whole class of mix-ups.

2. Compilation results are keyed by name. artifact_paths changes from List[Path] to Dict[str, Path], using the same keys as the on-device runner's pte_paths. The inference stage can then pick each artifact by name instead of relying on list order. Artifacts that weren't built are left out of the dict, so a check like "vision_encoder" in paths tells you whether the model is multimodal.

3. One typed unit per graph. GraphBundle replaces several parallel dicts with a single {graph_name: GraphBundle} map. Its checks run at construction time, so bad input fails early instead of as a KeyError during lowering.

4. Less copy-paste in compile specs. QnnCompileSpecBuilder replaces the HTP/GPU branch that llama.py repeats three times. It also sets the x86 emulator limits once in the constructor: no weight sharing and no shared buffer.

5. A bridge to the existing LLM code. ControlArgs lets the new pipeline drive the existing LLM components, which read settings off an argparse.Namespace. A test compares every default against the real llama.py parser. It's meant to be temporary.

6. Two real bugs fixed (details below).

This leaves Stream B ready for PR-B2 (multi-graph lowering, weight sharing, sharding, spill-fill). The llama.py flow is untouched.

---

Technical walkthrough

Area	What changed	Why it matters
artifact_keys.py / graph_names.py	String constants plus ALL_ARTIFACT_KEYS, DECODER_GRAPH_NAMES (decode first)	The package doesn't import from examples/. Tests check these strings still match decoder_constants.py.
graph_bundle.py	Frozen dataclass; __post_init__ requires quant_io_dtypes to have both kv_type and io_type, or be None	Catches a half-filled mapping when the bundle is built, not during lowering. The docstring is clear that meta and module can still be changed in place.
compile_spec_builder.py	resolve_soc_model builds its lookup from get_soc_to_chipset_map(); resolve_backend_type builds its from the enum's own __str__	No hand-written tables to fall out of date. Bad input raises ValueError listing the valid values, instead of an AttributeError from a failed getattr.
control_args.py	@dataclass that subclasses argparse.Namespace; build_parser, from_namespace, from_pipeline_context	The subclass is needed because QnnConfig.load_config checks isinstance(..., argparse.Namespace). extra_options can't override the four settings the context owns, so it can't silently switch the target SoC.
DefaultCompilerAdapter	Lowering with to_edge_transform_and_lower_to_qnn, then to_executorch with MemoryPlanningPass(alloc_graph_input/output=False) + BuildQuantIo, then writes the file	Same settings as the legacy flow. Before lowering, it decodes the target back out of compile_specs with flatbuffer_to_option and checks it against soc_model / backend_type, so a mismatch fails before compilation instead of on device.
Bug fix: push_artifacts	artifact_paths → artifact_paths.values()	Without this, the dict change would have pushed key strings like "text_decoder" to the device instead of file paths.
Bug fix: compile_spec import	compile_spec → compile_spec_schema	The old import pointed at a module that doesn't exist. It was hidden behind TYPE_CHECKING.

The layering is sound. The adapter compiles exactly one graph. Grouping several graphs into one .pte and spill-fill sizing are left to the strategy (PR-B2), which matches the rule set in Phase 1. I also like that soc_model and backend_type are now checked against the specs; before, they were accepted and then ignored.

---

Review findings

1. Once real specs are passed, the pipeline would always fail the target check (the SoC is a string, not the enum). default_compiler_adapter.py · _verify_target_matches_specs
GenAIPipeline._run_compilation passes soc_model=context.soc_model, which is a str like "SM8750" (genai_pipeline.py:269, with a comment saying the conversion is left to the adapter layer). The check then compares soc_model != options.soc_info.soc_model, and the right side is a QcomChipset (an IntEnum). "SM8750" != QcomChipset.SM8750 is always True, so it would raise ValueError even when the SoC is correct. Today nothing triggers this, because the pipeline doesn't pass compile_specs yet. That makes _decode_qnn_options return None and skip the check. It will trigger as soon as PR-B2 passes real specs. The tests only use enum values, so they don't catch it. Suggested fix: compare against resolve_soc_model(soc_model), and do the same with resolve_backend_type if backend_type can be a string. Add a test that passes a str SoC.
Fix this →

2. ControlArgs.build_parser() turns Optional[int] settings into strings. control_args.py · build_parser
The parser picks each argument's type from its default value. For fields that default to None (max_context_len, calib_num_fewshot, eval_num_fewshot, seed), that means type=str. llama.py declares these as type=int (for example llama.py:480-483), so --max-context-len 1024 produces "1024" instead of 1024. The code comment says "the consumer coerces", but the legacy code the comment refers to doesn't. It uses these values directly in arithmetic and comparisons, which would fail or behave oddly. The existing default-parity test doesn't catch this because it only compares defaults, not parsed values. Suggested fix: because of from __future__ import annotations, spec.type is just a string. Use typing.get_type_hints(cls) to get the real type, unwrap Optional[...], and use that as type=. Also add a test that parses a few values through both parsers and compares the results.
Fix this →

3. _decode_qnn_options swallows every exception (minor).
except Exception: return None is there so test stubs work. But it also means a real QNN spec that fails to decode, for example after a schema change, silently turns the target check off. I'd narrow this: skip entries that aren't QNN specs, but let a decode failure on an actual QCOM_QNN_COMPILE_SPEC entry raise, or at least log a warning.

4. executorch_config can be set in two places (minor).
GraphBundle.executorch_config and extra_options["executorch_backend_config"] both set the to_executorch config, but the adapter only reads the second. That's fine for now, since the adapter doesn't take bundles yet. PR-B2 should state which one wins so they don't conflict.

5. Docstring nits.

• CompilationInputConfig.graphs points to decoder_constants.DECODER_GRAPH_NAMES. It should point to this package's graph_names.DECODER_GRAPH_NAMES, since not depending on the example scripts is the whole point.
• GraphBundle refers several times to compilation/pass_policy.py, which doesn't exist until PR-B2. Worth marking that as upcoming.
• Style: the repo's CLAUDE.md asks for minimal comments, with comments only for non-obvious global context. Many of the docstrings here are long design notes, for example the GraphBundle section on the four fields it deliberately doesn't have, which mentions "removed … at review request". That kind of reasoning fits the PR description or design doc better than code that stays in the tree. Consider trimming.

6. CI. The one failure, unittest-editable / macos (torchaudio), looks unrelated to this PR.

---

Verdict

The foundation looks solid: clear naming, a reasonable decision on what is frozen, and the two bug fixes are real. I'd fix #1 before PR-B2 connects real compile specs to the adapter, and #2 before any CLI uses build_parser(). The rest are minor.

I didn't run the test suite: the sandbox didn't approve Python commands, so this review is from reading the code. Findings #1 and #2 are confirmed by tracing the call sites and the llama.py parser definitions.

[qti-horodnic]

Regarding Claude's comments:

1. Valid comment, fixed the bug and added a test which passes a string for SoC.
2. Valid comment, fixed the bug and added a test comparing parsed values through both parsers.
3. Fixed, a non-iterable stub still returns None, but a decode failure on a spec actually keyed QCOM_QNN_COMPILE_SPEC now raises rather than silently disabling the check.
4. executorch_config is ambiguous but harmless here since compile_model doesn't take bundles, so I'll state the precedence in the B2 that adds compile_graph_group, which does take in the bundle.
5. Regarding the docstrings, fixed the minor inaccuracies. Will keep them verbose for now until phase 2 is complete to help the reviewers and maintain synchronization between the developers. Once phase 2 is done, I'll trim the docstrings across the framework accordingly.

@psiddh Kindly take a look when you get chance, I think this PR should be GTG now.