Changes to SDPA to support no kv cache export

[tarun292]

Stack from ghstack (oldest at bottom):

• Add test case to export, quantize and lower vision encoder model for ET #7532
• Add utils to replace torchtune SDPA with ET Custom SDPA #7531
• -> Changes to SDPA to support no kv cache export #7530

Differential Revision: D67878163

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures

As of commit 84e7573 with merge base 68c0208 ():

NEW FAILURES - The following jobs have failed:

• Lint / lintrunner / linux-job (gh)
  >>> Lint for examples/models/llama/source_transformation/sdpa.py:
• pull / test-llama-runner-linux (fp32, xnnpack+custom+quantize_kv) / linux-job (gh)
  RuntimeError: Command docker exec -t 3c3a5cd191d7d2be0d350d5bbdcb9bff7bc3a49c05cf2714f0ecd7de20a4b2b1 /exec failed with exit code 1

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

[tarun292]

@tarun292 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[dvorjackz]

Tbh I think you can just set this to 0, should work for the no kv cache text decoder as well since it represents start position, so you don't need to set input_pos to torch.tensor(0) in your other pr

[tarun292]

@tarun292 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[tarun292]

@tarun292 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.