Add Initial Compile for Llama 3.2 11B: Decoder TransformerSelfAttentionLayer, TransformerCrossAttentionLayer

[Jack-Khuu]

Based on https://github.com/pytorch/torchtune/blob/57ab583c84c4a9dcacac23aeabc81f2a679670fe/torchtune/training/_compile.py#L42-L52

Compile as much of the Llama3.2 11B Vision Model from tune as we can

---

Testing:

 python torchchat.py generate llama3.2-11B --prompt "What's in this image?" --image-prompt assets/dog.jpg --compile

Run it a few times to avoid caching times: On the 4th iteration

What's in this image?The image is of a dog wearing sunglasses and riding a skateboard in the street. The dog is white with some brown spots on its face and ears, and it has its tongue out. It's wearing red sunglasses and a blue collar, and its front paws are on a red penny board with yellow wheels. The background is a road with green trees on either side.2024-10-09:18:43:33,307 INFO     [generate.py:1158] 
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~                
Generated 74 tokens                 
Time for inference 3: 1.9990 sec total                 
Time to first token: 0.1759 sec with parallel prefill.                

      Total throughput: 37.5180 tokens/sec, 0.0267 s/token                 
First token throughput: 5.6850 tokens/sec, 0.1759 s/token                 
 Next token throughput: 40.5893 tokens/sec, 0.0246 s/token                     
2024-10-09:18:43:33,307 INFO     [generate.py:1169] 
Bandwidth achieved: 798.83 GB/s

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchchat/1287

• 📄 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 1b83644 with merge base 438ebb1 ():
💚 Looks good so far! There are no failures yet. 💚

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

[ebsmothers]

Out of curiosity, how long does it take to compile in this case?

[Jack-Khuu]

About 2 Minutes for Llama 3.1

[iseeyuan]

LGTM. Thanks!