[Executorch][llama] Refactoring sdpa

[kimishpatel]

Stack from ghstack (oldest at bottom):

• [Executorch][quant] Optimize per channel dequantize #5656
• [Executorch][llama] Add custom_sdpa and use that instead of sdpa_with_kv_cache #5655
• [Executorch] Update EXECUTORCH_LIBRARY macro #5654
• -> [Executorch][llama] Refactoring sdpa #5653
• [Executorch][llama] Update SDPA op to use quantized kv cache #5652
• Add update_quantized_cache op #5527
• Refactor custom SDPA op to separate kv cache update from the custom sdpa op #5651
• [ExecuTorch] Add quantized kv cache to llama #5650
• Fix dequantize per channel to handle double scale type #5524
• [ExecuTorch] Some updated to kv cache #5649

• Change variable name from is_sdpa_with_kv_cache to is_seq_at_dim_1 to be more meaningful
• remove seq_len as it can be derived from the size of query

Differential Revision: D62623243

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 6 New Failures

As of commit a80ec36 with merge base d2ba238 ():

NEW FAILURES - The following jobs have failed:

• pull / test-llama-runner-linux (bf16, buck2, portable) / linux-job (gh)
  RuntimeError: Command docker exec -t a30a8398f9be992ecf4ba2a8e6241c36e76df2ab1d4170096375dc925e099230 /exec failed with exit code 3
• pull / test-llama-runner-linux (fp32, buck2, portable) / linux-job (gh)
  RuntimeError: Command docker exec -t 2aae7c342fed4390e64926adcc8fb5e6c0df774d46c8cd128d722d810d9dc72e /exec failed with exit code 3
• pull / test-llama-runner-linux (fp32, buck2, xnnpack+custom) / linux-job (gh)
  RuntimeError: Command docker exec -t ca69dd72258c7bebb9bd6baa4b934d750c8de09850117f9236121a8f659966a6 /exec failed with exit code 3
• pull / test-llama-runner-linux (fp32, buck2, xnnpack+custom+qe) / linux-job (gh)
  RuntimeError: Command docker exec -t 25dd0132a0d4d3b516b35cfcf8a094441dddebf360169f2792d8d0d29718c399 /exec failed with exit code 3
• pull / unittest / linux / linux-job (gh)
  kernels/quantized/test/test_quant_dequant_per_token.py::QuantizePerTokenTest::test_quantize_per_token_dynamic
• pull / unittest / macos / macos-job (gh)
  kernels/quantized/test/test_quant_dequant_per_token.py::QuantizePerTokenTest::test_quantize_per_token_high_rank

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: D62623243

[facebook-github-bot]

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