Skip to content

Export hf_static float-cache programs and register their Python runner - #24571

Draft
psiddh wants to merge 6 commits into
pytorch:mainfrom
psiddh:stack2/pr06-hf-static-export
Draft

psiddh wants to merge 6 commits into
pytorch:mainfrom
psiddh:stack2/pr06-hf-static-export

Conversation

@psiddh

@psiddh psiddh commented Nov 3, 2026 •

Copy link
Copy Markdown
Contributor

Export hf_static float-cache programs and register their Python runner

HF exports a variable-length static-cache graph that writes into caller-owned buffers. HTP needs fixed shapes and QNN cannot perform those input mutations. Export fixed prefill/decode methods, attend over incoming caches and current-token rows, and retain cache writes, masks and counters in CPU kernels with the existing reinplace pass.

Each step keeps HF's tensor interface: token IDs, masks, positions and full K/V caches in; logits and those caches updated in place out. The fixed-shape padding/counter adaptation is hf_static v1. Decode is required; prefill is optional. Caches remain float16/float32 even with quantized compute. Quantized caches, sha and full-length sliding mode are excluded.

Backend-neutral rewriting and cache_graph.py stay in the LLM extension. Tensor dtype and constant-method decoding live in their sole driver module. QNN partitioning, SeqMSE, shared activation encodings, calibration prompts and the hf_static-only CLI live in backends/qualcomm/utils/hf_static.py. Their retained helper implementations and constants are unchanged. The shared Gemma folding and Qwen3 norm-layout tests remain covered here.

The guarded transformers hook transformers.exporters.register_io_contract is available only in the local transformers branch used for this work, not stock transformers. Without the hook, registration is skipped. The Python driver handles padding, absolute counters, reset, context-tail decode fallback and DLPack logits; it rejects planned/device-resident state. No zero-copy Python host API is claimed. External cache layout requires export-time validation.

Review shared graph utilities and rewriting, Python driver/registration, backend lowering/tuning and the CLI, then tests. The C++ manager follows in PR 7. Cache inspection, adopt/fork/evict policies remain IOManager follow-ups.

Files

backends/qualcomm/_passes/convert_linear_to_conv2d.py
backends/qualcomm/quantizer/custom_annotation.py
backends/qualcomm/quantizer/hf_llm_quant_recipe.py
backends/qualcomm/tests/test_hf_static.py
backends/qualcomm/utils/hf_static.py
backends/qualcomm/utils/targets.bzl
extension/llm/export/BUCK
extension/llm/export/cache_graph.py
extension/llm/export/hf_static.py
extension/llm/export/test/BUCK
extension/llm/export/test/test_hf_static.py
extension/llm/io_contracts/BUCK
extension/llm/io_contracts/__init__.py
extension/llm/io_contracts/_transformers.py
extension/llm/io_contracts/hf_static.py
extension/llm/io_contracts/test/__init__.py
extension/llm/io_contracts/test/test_hf_static.py

Test Plan

Commit 7cb492766bdd019340a865774dc8742489f10c6b on 892e50425f5a3e90ff8d3c8802fdfb31f8599518. Validated locally with Python 3.12, pinned Torch/TorchAO, the local transformers exporter branch and QAIRT SDK 2.48.

['flock', '/data/users/sidart/qnnflow/gh_prs/model-export.lock', '/data/users/sidart/qnnflow/gh_prs/venv-rebase312/bin/python', '-c', 'import torch,pytest,sys;torch.ops.load_library("/data/users/sidart/qnnflow/gh_prs/build-rebase-native-exir-test/libexir_test.so");raise SystemExit(pytest.main(sys.argv[1:]))', '-q', '-rs', 'backends/qualcomm/tests/test_hf_static.py', 'extension/llm/export/test/test_hf_static.py', 'extension/llm/io_contracts/test/test_hf_static.py', 'exir/tests/test_partitioner_transform_passes.py', 'backends/qualcomm/tests/test_output_order.py', 'backends/qualcomm/tests/test_partitioner_options.py', '--junitxml=/data/users/sidart/qnnflow/gh_prs/logs/split-pr06-python-1.xml']
97 passed, 55 subtests passed in 342.89s (0:05:42)

DLPack/non-Torch output subprocess: 3 passed + 2 subtests, zero failures. CLI help and forwarding checks pass. All 32 retained function/class definitions are structurally equivalent, with identical moved constants and calibration prompts.

No device/physical rpcmem, two-target FCB, real-model export or full QNN device-suite result is claimed. Buck could not start because the local launcher requires a missing .buck2. The base already contains #23563 and #23183; neither is a separate prerequisite commit here.

Related PRs (review in order)

  1. Partitioner: run backend-contributed edge transform passes #24566 Partitioner: run backend-contributed edge transform passes
  2. Export: avoid reused lifted-constant names #24567 Export: avoid reused lifted-constant names
  3. LLM runner: select I/O contracts and accept custom buffer allocators #24568 LLM runner: select I/O contracts and accept custom buffer allocators
  4. QNN: construct partitioners from options and allocate LLM I/O in rpcmem #24569 QNN: construct partitioners from options and allocate LLM I/O in rpcmem
  5. QNN: bind delegate outputs in exported program order #24570 QNN: bind delegate outputs in exported program order
  6. Export hf_static float-cache programs and register their Python runner #24571 Export hf_static float-cache programs and register their Python runner
  7. LLM runner: validate and drive hf_static float-cache programs #24572 LLM runner: validate and drive hf_static float-cache programs
  8. Document the version-1 hf_static I/O contract #24573 Document the version-1 hf_static I/O contract

Each PR includes the commits of the ones above it; review only its own top commit.

@pytorch-bot

pytorch-bot Bot commented Nov 3, 2026 •

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❌ 2 New Failures, 1 Cancelled Job, 1 Unrelated Failure, 1 Unclassified Failure

As of commit 3ce285a with merge base 765e973 (image):

NEW FAILURES - The following jobs have failed:

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

  • pull / test-arm-backend-docgen / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
    RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1

CANCELLED JOB - The following job was cancelled. Please retry:

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

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

@github-actions

github-actions Bot commented Nov 3, 2026

Copy link
Copy Markdown

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

Run each partitioner’s edge transforms after caller-supplied passes so backends can use the standard lowering entry point. QNN contributes its existing transforms through the new hook.
Reserve names for constant tensors already present in the exported program before lifting new constants, so methods with folded tables retain distinct constants.
Choose the registered IOManager and its logits, method names and step sizes, rejecting unknown contracts and invalid metadata. Load and reset the selected methods; route prompt tails to decode with a side-effect-free capacity check. Preserve module and I/O-buffer lifetimes and expose a backend-neutral allocator hook.
Route backend options through QnnPartitioner.from_options and register an allocator that supplies QNN shared-memory buffers through the generic IOManager hook. The LLM runner remains backend-neutral.
Encode graph output positions and restore their order at both runtime tensor-allocation entry points. Preserve context-binary order and legacy mutable outputs, and ignore malformed or overflowing position labels without throwing. This runtime change needs Qualcomm review.
Move backend-neutral static-cache rewriting into the LLM extension and keep QNN partitioning and tuning in the backend. Detect sliding layers from model configuration, keep cache writes in place with standard reinplace, and register a versioned runner that handles fixed chunks, padding, counters and reset. Quantized caches, sha attention and full-length sliding mode are excluded.
@psiddh
psiddh force-pushed the stack2/pr06-hf-static-export branch from 82bb0a1 to 3ce285a Compare November 13, 2026 08:50

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant