Add customized static cache implementation

[helunwencser]

This PR is based on the prototype done in huggingface/transformers#31706. We want to export HF models to ExecuTorch. However, we cannot do that right now due to its Cache is not a torch.nn.Module.

This PR tries to solve this problem by implementing a customized StaticCache which is not only a Cache but also a torch.nn.Module. Most of its implementation is copied from transformers with the following modification:

• a subclass of torch.nn.Module
• add register_buffer call, copied from [Demo][ExecuTorch] Lower and run native Gemma e2e in ExecuTorch huggingface/transformers#31706
• fix two bugs in the StaticCache implementation: 1. get_seq_length should return a number instead of tensor, 2. update should only returned filled cache slots instead of the whole static cache.

Test Plan:
Make sure the following commands generate the exact same output but the later one, with kv cache enabled, is faster:

python3 -m examples.models.phi-3-mini.eager
python3 -m examples.models.phi-3-mini.eager -kv

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit e31ff36 with merge base 28cfabb ():
💚 Looks good so far! There are no failures yet. 💚

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