Update KV Cache to use num_kv_heads instead of num_heads

[mirceamironenco]

Context

What is the purpose of this PR? Is it to

• add a new feature
• fix a bug
• update tests and/or documentation
• other (please add here)

Please link to any issues this PR addresses.

Changelog

What are the changes made in this PR?

• Update KV Cache to use num_kv_heads instead of num_heads for more memory-efficient generation.
• Simplified/adapted mha fwd pass to have the expand happen after the kv cache.

Test plan

Please make sure to do each of the following if applicable to your PR. If you're unsure about any one of these just ask and we will happily help. We also have a contributing page for some guidance on contributing.

• run pre-commit hooks and linters (make sure you've first installed via pre-commit install)
• add unit tests for any new functionality
• update docstrings for any new or updated methods or classes
• run unit tests via pytest tests
• run recipe tests via pytest tests -m integration_test
• manually run any new or modified recipes with sufficient proof of correctness
• include relevant commands and any other artifacts in this summary (pastes of loss curves, eval results, etc.)

UX

If your function changed a public API, please add a dummy example of what the user experience will look like when calling it.
Here is a docstring example
and a tutorial example

• I did not change any public API
• I have added an example to docs or docstrings

If one changes generation.yaml to use llama 3.1 8b (num_kv_heads < num_heads), and the prompt/max_new_toks:

INFO:torchtune.utils._logging:Running InferenceRecipe with resolved config:

checkpointer:
  _component_: torchtune.training.FullModelHFCheckpointer
  checkpoint_dir: /tmp/Meta-Llama-3.1-8B-Instruct/
  checkpoint_files:
  - model-00001-of-00004.safetensors
  - model-00002-of-00004.safetensors
  - model-00003-of-00004.safetensors
  - model-00004-of-00004.safetensors
  model_type: LLAMA3
  output_dir: /tmp/Meta-Llama-3.1-8B-Instruct/
  recipe_checkpoint: null
device: cuda
dtype: bf16
enable_kv_cache: true
max_new_tokens: 4096
model:
  _component_: torchtune.models.llama3_1.llama3_1_8b
prompt:
  system: null
  user: Tell me a long joke.
quantizer: null
seed: 1234
temperature: 0.6
tokenizer:
  _component_: torchtune.models.llama3.llama3_tokenizer
  max_seq_len: null
  path: /tmp/Meta-Llama-3.1-8B-Instruct/original/tokenizer.model
  prompt_template: null
top_k: 300

I get with the new kv cache (on RTX A6000):

INFO:torchtune.utils._logging:Time for inference: 12.72 sec total, 22.79 tokens/sec
INFO:torchtune.utils._logging:Bandwidth achieved: 379.93 GB/s
INFO:torchtune.utils._logging:Memory used: 17.01 GB

Old kv cache:

INFO:torchtune.utils._logging:Time for inference: 11.86 sec total, 24.46 tokens/sec
INFO:torchtune.utils._logging:Bandwidth achieved: 447.16 GB/s
INFO:torchtune.utils._logging:Memory used: 18.62 GB

I'm not sure how relevant the tok/sec change is given that this is not compiled and batch_size=1.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/1961

• 📄 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 5f0d2b7 with merge base 08efaed ():
💚 Looks good so far! There are no failures yet. 💚

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

[SalmanMohammadi]

I like this

[mirceamironenco]

Addressed these and also updated the generate benchmark for single device.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 67.26%. Comparing base (9eced21) to head (65f7498).
Report is 13 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1961      +/-   ##
==========================================
- Coverage   68.40%   67.26%   -1.14%     
==========================================
  Files         311      316       +5     
  Lines       16973    17342     +369     
==========================================
+ Hits        11610    11665      +55     
- Misses       5363     5677     +314     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[SalmanMohammadi]

Thanks so much for this, and for your patience in helping test it.