[Executorch][llama] Add custom_sdpa and use that instead of sdpa_with_kv_cache

[kimishpatel]

Stack from ghstack (oldest at bottom):

• [Executorch][quant] Optimize per channel dequantize #5622
• -> [Executorch][llama] Add custom_sdpa and use that instead of sdpa_with_kv_cache #5621
• [Executorch] Update EXECUTORCH_LIBRARY macro #5620
• [Executorch][llama] Refactoring sdpa #5619
• [Executorch][llama] Update SDPA op to use quantized kv cache #5618
• Add update_quantized_cache op #5527
• Refactor custom SDPA op to separate kv cache update from the custom sdpa op #5617
• [ExecuTorch] Add quantized kv cache to llama #5616
• Fix dequantize per channel to handle double scale type #5524
• [ExecuTorch] Some updated to kv cache #5615

sdpa_with_kv_cache updates kv cache. In quantized kv cache, cache updates
happens separately. Then the quantized cache is dequantized. After that
we call sdpa_with_kv_cache which copies k and v data into dequantized cache.
Although this is not needed because the actual cache is the one that is
quantized.
For very large context length this will add significant amount data copy.

Subsequent diffs will deprecate sdpa_with_kv_cache op and deconstruct that
using a) update_cache op and b) custom_sdpa op.

Differential Revision: D62623241

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 540e91e with merge base b2517d6 ():

NEW FAILURES - The following jobs have failed:

• pull / test-llama-runner-linux (bf16, buck2, portable) / linux-job (gh)
  RuntimeError: Command docker exec -t b300d4d8be66703458f0c412bc402e359c397de7b4e6d4a4fe590c8c2ca7e780 /exec failed with exit code 3
• pull / test-llama-runner-linux (fp32, buck2, portable) / linux-job (gh)
  RuntimeError: Command docker exec -t 3f34142f852f2d258c13eae844845d4ab25b536eb0f96fb2ef5e9c1218b5a9dd /exec failed with exit code 3
• pull / test-llama-runner-linux (fp32, buck2, xnnpack+custom) / linux-job (gh)
  RuntimeError: Command docker exec -t bbfa710d7ffae46945c59eacd2f00b425e463aa63cbf316335213f84d8ccb0ec /exec failed with exit code 3
• pull / test-llama-runner-linux (fp32, buck2, xnnpack+custom+qe) / linux-job (gh)
  RuntimeError: Command docker exec -t b3225413a7a55d28483235f5896bf15c9d84b1449315d6f49f07d7631aaf7a23 /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_large_tensor
• pull / unittest / macos / macos-job (gh)
  kernels/quantized/test/test_quant_dequant_per_token.py::QuantizePerTokenTest::test_quantize_per_token_large_tensor

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