[Attention] Support FP4 attention quantization

[kylesayrs]

Purpose

• Support FP4 attention quantization

Changes

• Add tensor group option to _flatten_attention
• Update test_static_attention_quantization to use a larger shape

Testing

• Added tests
• https://huggingface.co/kylesayrs/Llama-3.2-1B-Instruct-attention-nvfp4-head

[gemini-code-assist]

Summary of Changes

Hello @kylesayrs, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enhances the quantization capabilities by integrating FP4 attention quantization, particularly for the TENSOR_GROUP strategy. This change is vital for optimizing the memory and computational efficiency of large language models by allowing more granular control over attention mechanism quantization, which can lead to faster inference and reduced resource consumption.

Highlights

• FP4 Attention Quantization Support: This pull request introduces support for FP4 attention quantization, specifically enabling the TENSOR_GROUP strategy for attention mechanisms within the _flatten_attention helper function.
• _flatten_attention Functionality Update: The _flatten_attention function in src/llmcompressor/observers/helpers.py has been updated to correctly handle the TENSOR_GROUP quantization strategy, allowing attention tensors to be reshaped for group-based quantization instead of raising an error.
• New Test Cases for Quantization: New test cases have been added to tests/llmcompressor/modifiers/calibration/test_lifecycle.py to validate the FP4 tensor_group attention quantization, ensuring its correct functionality and expected loss. Existing test parameters were also adjusted.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces support for FP4 attention quantization by implementing the necessary tensor flattening logic for group-based quantization strategies. The changes are well-supported by the addition of a new test case and updates to existing tests, which now use a more standard tensor dimension order for attention. The implementation appears correct and is a good addition. I have one minor suggestion to improve code clarity by fixing a typo in a comment.