[Transform] Spinquant R3

[kylesayrs]

Purpose

• Support R3 for spinquant (without attention quantization, but possibly with kv_cache quantization)

Prerequisites

• [Transform] SpinQuant R4 #1746
• [Transform] Attention/Cache transforms

Changes

• Break out head dim inference
• Add attention target for spinquant mappings

Testing

• Tested model coherence with llama 8b model

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[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 enables R3 rotation support for Spinquant, primarily focusing on KV-cache quantization. It achieves this by refactoring the head dimension inference logic into a dedicated utility method and introducing a new attention target within the Spinquant mapping system. The changes facilitate the application of R3 schemes to relevant attention components, and the coherence of a Llama 8B model was tested to validate the implementation.

Highlights

• Spinquant R3 Support: This PR introduces support for the R3 rotation scheme within Spinquant, specifically enabling KV-cache quantization without attention quantization.
• Head Dimension Inference Refactoring: The logic for inferring the head dimension has been extracted into a dedicated private method (_infer_head_dim) to improve code organization and reusability across different rotation schemes.
• Attention Target Mapping: A new attn target has been added to the SpinQuantMapping to specifically target the self_attn module, which is utilized by the new R3 scheme for q_attn and k_cache locations.

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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 adds support for SpinQuant R3 rotations and also implements R4. The refactoring to centralize head dimension inference is a good improvement. I've found a critical issue in the implementation of the R4 scheme where it seems to be applying an identity transformation, making it ineffective. I've also suggested a minor improvement to an error message for better debuggability. Please see the detailed comments.

[brian-dellabetta]

looks good!

The base branch was changed.