Swap mha (#6719)

* Swap mha

Move to extension/llm/modules

Lint

Add tests

* Fix tests

* Delete old file

Author: jackzhxng

extension/llm/modules/README.md

@@ -1,14 +1,17 @@
-## Export Friendly Modules
+## Export-friendly Modules
 
-Modules in this directory are:
-* Extending `torch.nn.Module`.
-* Guranteed to work out of the box with `torch.export.export()` and `torch.aot_compile()`.
-* Guranteed to be able to work with ExecuTorch.
+Modules in this directory:
+* Extend `torch.nn.Module`.
+* Are guaranteed to work out of the box with `torch.export.export()`.
+* Should work out of the box with `torch.aot_compile()`.
+* Should be able to workt with ExecuTorch.
 
 All modules should be covered by unit tests to make sure they are:
-1. giving the same output as the reference implementation in PyTorch or torchtune
-2. export friendly
-3. AOTI friendly
-4. ExecuTorch friendly
+1. Give the output as the reference eager model in PyTorch or TorrchTune
+2. Export-friendly
 
-Notice that these modules are subject to change (may upstream to torchtune) so proceed with caution.
+Additionally, we aim to make these modules:
+3. AOTI-friendly
+4. ExecuTorch-friendly
+
+These modules are subject to change (may upstream to TorchTune) so proceed with caution.