[Executorch][llm] Enable local global attention in export_llama script

[kimishpatel]

Stack from ghstack (oldest at bottom):

• -> [Executorch][llm] Enable local global attention in export_llama script #10612
• [Executorch][llm] Enable leveraging ring kv cache via module swap #10611
• [Executorch][llm] Make custom update cache op operate on indices #10610
• [Executorch][llm] Add ring buffer based kv cache and mask calculation to MHA #10609
• [Executorch][llm] Add support for ring kv cache and ring attention #10608

Added a new option of --local_global_attention that takes in pattern of sizes to determine which layers are using local sliding window attention.
For example, [0, 256, 256, 0, 256, 256] can be used for 6 layers transformer. Or you can also use [0, 256, 256] as pattern you want
to repeat.

Differential Revision: D73891423

cc @larryliu0820 @mergennachin @cccclai @helunwencser @jackzhxng

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• CI workflows being skipped on PR

❌ 6 New Failures

As of commit 678ba9f with merge base cd3b53d ():

NEW FAILURES - The following jobs have failed:

• pull / test-llama-runner-qnn-linux (fp32, qnn_16a16w, qnn) / linux-job (gh)
  RuntimeError: Command docker exec -t 0593698797f823406e7152eff851ae7013c6c28e85471f329b7bf9b25f514908 /exec failed with exit code 1
• pull / test-llama-runner-qnn-linux (fp32, qnn_8a8w, qnn) / linux-job (gh)
  RuntimeError: Command docker exec -t dbc3df9b9418598d7d22b45fd6230fdc54ffed97e43981565dff5ebc0148903c /exec failed with exit code 1
• pull / unittest / linux / linux-job (gh)
  examples/models/llama/tests/test_simple_sdpa.py::SDPATest::test_simple_sdpa
• pull / unittest / macos / macos-job (gh)
  backends/xnnpack/test/ops/test_conv1d.py::TestConv1d::test_qs8_conv1d_batchnorm_seq
• pull / unittest-editable / linux / linux-job (gh)
  examples/models/llama/tests/test_simple_sdpa.py::SDPATest::test_simple_sdpa
• pull / unittest-editable / macos / macos-job (gh)
  backends/xnnpack/test/ops/test_sqrt.py::TestSqrt::test_fp16_sqrt

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

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D73891423

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D73891423

[jackzhxng]

From my understanding local global attention is tied closely to the model, and isn't something you can adjust given a particular checkpoint unlike these other export options such as qmode, enable_kv_cache, etc. I think it would be better to have this in model_args instead so that we can represent it in the config json file, so it is part of model configuration, not export configuration

[kimishpatel]

From my understanding local global attention is tied closely to the model, and isn't something you can adjust given a particular checkpoint unlike these other export options such as qmode, enable_kv_cache, etc. I think it would be better to have this in model_args instead so that we can represent it in the config json file, so it is part of model configuration, not export configuration

You can also make all the layers just do sliding window attention for example, so you can view this also as pure optimization.

I havent thought a lot of about making this a model arg vs. export arg. I think your point definitely makes sense but I think it can be configurable as well if you want to support infinite generation as many runners do.