Skip to content

Harden custom SDPA export path - #25032

Open
bdemirb wants to merge 2 commits into
pytorch:mainfrom
bdemirb:tmp-1366005-rebase2
Open

bdemirb wants to merge 2 commits into
pytorch:mainfrom
bdemirb:tmp-1366005-rebase2

Conversation

@bdemirb

@bdemirb bdemirb commented Nov 13, 2026 •

Copy link
Copy Markdown
Collaborator

Tighten the LLM custom SDPA export rewrite so it can be exercised in an export-only environment before any Arm backend lowering is added. The pass now accepts GQA decode shapes only when the query and KV head layout is compatible, and it cleanly falls back to the Python llama op registration when the compiled custom-op library is unavailable.

Add export-pass coverage for a representative decode-time GQA shape and for the unsupported head-mismatch case without enable_gqa. The fallback registration also provides a Python reference body for llama.custom_sdpa so the rewritten graph remains executable in local tests.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani

@pytorch-bot

pytorch-bot Bot commented Nov 13, 2026 •

Copy link
Copy Markdown

🔗 Helpful Links

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

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

⏳ 2 Pending, 1 Unrelated Failure

As of commit 6dd073e with merge base 8c77fe1 (image):

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.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Nov 13, 2026
@bdemirb

bdemirb commented Nov 13, 2026

Copy link
Copy Markdown
Collaborator Author

@pytorchbot label "partner: arm"

@pytorch-bot pytorch-bot Bot added the partner: arm For backend delegation, kernels, demo, etc. from the 3rd-party partner, Arm label Nov 13, 2026
@bdemirb

bdemirb commented Nov 13, 2026

Copy link
Copy Markdown
Collaborator Author

@pytorchbot label "release notes: arm"

@pytorch-bot pytorch-bot Bot added the release notes: arm Changes to the ARM backend delegate label Nov 13, 2026
Tighten the LLM custom SDPA export rewrite so it can be exercised in an
export-only environment before any Arm backend lowering is added. The
pass now accepts GQA decode shapes only when the query and KV head
layout is compatible, and it cleanly falls back to the Python llama op
registration when the compiled custom-op library is unavailable.

Add export-pass coverage for a representative decode-time GQA shape and
for the unsupported head-mismatch case without enable_gqa. The fallback
registration also provides a Python reference body for llama.custom_sdpa
so the rewritten graph remains executable in local tests.

Test:
PYTHONPATH=/private/tmp/executorch_patch1_sdpa/src \
PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \
/Users/bardem01/devkit7/executorch/.venv-fvp/bin/pytest -q \
extension/llm/export/test_export_passes.py

Authored with assistance from OpenAI Codex.

Signed-off-by: Baris Demir <baris.demir@arm.com>
Change-Id: I880b2dbd201218467254db2eb5a981f00254623b
@bdemirb
bdemirb force-pushed the tmp-1366005-rebase2 branch from 07f0c96 to 154f112 Compare November 13, 2026 11:27
@bdemirb bdemirb changed the title Arm backend: Harden custom SDPA export path Harden custom SDPA export path Nov 13, 2026
@zingo zingo added the help wanted Extra attention is needed label Nov 13, 2026

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. help wanted Extra attention is needed partner: arm For backend delegation, kernels, demo, etc. from the 3rd-party partner, Arm release notes: arm Changes to the ARM backend delegate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants