add: GQA support to AWQ

[soyr-redhat]

SUMMARY:
Addresses #2808

AWQ currently skips v_proj -> o_proj smoothing for GQA models because _check_layers_are_compatible() rejects mappings where v_proj.out_features != o_proj.in_features. In GQA models (Llama 3, Mistral, Qwen2, etc.), this mismatch is expected, KV heads are repeated via repeat_kv before feeding into o_proj.

This PR adds GQA support by:

• Detecting valid GQA dimension ratios in the compatibility check instead of rejecting them
• Compressing scales from o_proj's input dimension to v_proj's output dimension at application time, by averaging across KV head repeat groups to match the repeat_kv layout
• Leaving the grid search and activation hooks unchanged, scales are computed at (hidden_dim,) for o_proj and only compressed when applying to v_proj

Fused qkv_proj with GQA is still lacking support. I'm unsure as to if this can be addressed by just knowing model dimensions.

Currently with runs, the logger outputs:

GQA detected for model.layers.0.self_attn.v_proj: num_repeats=n_0, head_dim=n_1

For every GQA layer.

TEST PLAN:

• All 22 existing AWQ unit tests pass
• 4 new GQA-specific unit tests pass ----- Attribution to Claude Code
• End-to-end AWQ on nm-testing/tinysmokellama-3.2 (GQA, num_repeats=2, head_dim=4)
• End-to-end AWQ on meta-llama/Llama-3.2-1B (GQA, num_repeats=4, head_dim=64): RTX 3090, ~42s

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 425ae012-6021-4724-841c-131f33953654

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

AWQModifier now detects and supports Grouped Query Attention patterns by extracting optional head dimension from model config, updating layer compatibility checks to return GQA metadata, storing this metadata in ResolvedMapping, and applying conditional scale compression during smoothing.

Changes

GQA Support in AWQ Modifier

Layer / File(s)	Summary
Data model: GQA metadata in ResolvedMapping src/llmcompressor/modifiers/transform/awq/mappings.py	ResolvedMapping dataclass adds optional gqa_head_dim field to store detected GQA head dimension metadata.
Head dimension extraction and GQA detection src/llmcompressor/modifiers/transform/awq/base.py	New _get_head_dim() helper extracts head dimension from model config, and _check_layers_are_compatible() is rewritten to accept head_dim parameter and return (compatible, gqa_head_dim) tuple with GQA validation logic. Mapping resolution captures detected GQA head dimension and logs detection info when present.
Scale compression for GQA and smoothing integration src/llmcompressor/modifiers/transform/awq/base.py	New _compress_scales_for_gqa() averages scales across GQA repeat groups. During smoothing, compressed scales replace raw scales when GQA is detected, applied to both 1D and multi-dimensional weight scaling and bias.
Unit tests for GQA detection and scale compression tests/llmcompressor/modifiers/transform/awq/test_base.py	Tests verify _set_resolved_mappings() accepts v_proj→o_proj mappings for GQA configurations, skips them when config absent, rejects fused qkv_proj patterns, and validate correctness of scale averaging in _compress_scales_for_gqa().

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related issues

• AWQ GQA support #2808: Implements explicit GQA support in AWQModifier with head dimension extraction, gqa_head_dim storage in ResolvedMapping, and scale-averaging behavior requested in that issue.

Suggested labels

awq, transforms, two-reviews

Suggested reviewers

• shanjiaz
• HDCharles
• brian-dellabetta

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 35.71% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'add: GQA support to AWQ' clearly and accurately summarizes the main change: adding Group Query Attention support to the AWQ quantization modifier.
Description check	✅ Passed	The description is detailed and directly related to the changeset, explaining the problem (AWQ skipping GQA models), the solution (detecting GQA patterns and compressing scales), testing approach, and known limitations.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

Warning

Review ran into problems

🔥 Problems

Linked repositories: Your configuration references 1 linked repositories, but your current plan allows 0. Analyzed ``, skipped vllm-project/compressed-tensors.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[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]

Code Review

This pull request adds support for Grouped Query Attention (GQA) in the AWQ modifier by extracting the head dimension from the model configuration, validating GQA ratios during compatibility checks, and compressing scales across GQA repeat groups during smoothing. Unit tests are also added to verify these changes. The review feedback highlights opportunities to improve robustness by replacing .view() with .reshape() to avoid RuntimeError on non-contiguous tensors, and adding a check to prevent a potential ZeroDivisionError if head_dim is zero.

[mergify]

Merge Protections

🔴 2 of 2 protections blocking · waiting on 👀 reviews

	Protection	Waiting on
🔴	Require one maintainer review	👀 reviews
🔴	Require two reviews	👀 reviews

🔴 Require one maintainer review

Waiting for

• #changes-requested-reviews-by = 0
• any of:
  • approved-reviews-by=HDCharles
  • approved-reviews-by=brian-dellabetta
  • approved-reviews-by=dsikka
  • approved-reviews-by=kylesayrs
  • approved-reviews-by=yiliu30

This rule is failing.

All PRs must have at least one approving review from a maintainer before merging.

• #changes-requested-reviews-by = 0
• any of:
  • approved-reviews-by=HDCharles
  • approved-reviews-by=brian-dellabetta
  • approved-reviews-by=dsikka
  • approved-reviews-by=kylesayrs
  • approved-reviews-by=yiliu30

🔴 Require two reviews

Waiting for

• #approved-reviews-by >= 2
• #changes-requested-reviews-by = 0

This rule is failing.

PRs labelled "two-reviews" must have at least two approving reviews before merging.

• #approved-reviews-by >= 2
• #changes-requested-reviews-by = 0

[soyr-redhat]

Ran a quick perplexity check on Llama 3.2 1B (wikitext-2-raw-v1, 2048 token window, 32 calibration samples):

Config	Perplexity
FP16 baseline	7.25
AWQ W4A16 without v_proj->o_proj smoothing	9.07
AWQ W4A16 with GQA smoothing	9.21

The results seem to be comparable, but with a 1B model, its not saying too much, will get a deployment running on h200s to test a larger model, but a 1.5% variance seems convincing.

[soyr-redhat]

Results (acc_norm where available, acc for winogrande)

Model	GQA Ratio	Config	arc_easy	hellaswag	piqa	winogrande
Llama 3.1 8B	4:1	FP16	0.825	0.793	0.812	0.745
		AWQ (upstream)	0.789	0.783	0.801	0.726
		AWQ + GQA smooth	0.790	0.787	0.796	0.710
Mistral 7B v0.3	4:1	FP16	0.802	0.807	0.821	0.741
		AWQ (upstream)	0.789	0.799	0.813	0.729
		AWQ + GQA smooth	0.789	0.802	0.819	0.731
Gemma 2 9B	2:1	FP16	0.878	0.800	0.829	0.741
		AWQ (upstream)	0.871	0.789	0.826	0.747
		AWQ + GQA smooth	0.873	0.793	0.826	0.740

Summary

The implemented GQA smoothing strategy seems appears to be quality-neutral across all three architectures and both GQA ratios (4:1 and 2:1).

All differences between upstream and GQA smooth are within a good range of eachother. The new v_proj->o_proj smoothing path does not degrade model quality.

[soyr-redhat]

Reviewed Gemini Code Review

[brian-dellabetta]

Hi @soyr-redhat , thanks for setting this up. The changes make sense, lemme get @HDCharles 's thoughts on the complexity vs improvement tradeoff here.

You expect this to be generally applicable to any GQA model, yes?

[soyr-redhat]

Hey @brian-dellabetta! Thanks for taking a look.

I personally do expect this to be generally applicable, as I ran some tests across a couple different architectures/ratios. At least as far as I'm aware, this should be agnostic since we're directly deriving the ratio from the model config!

[HDCharles]

i think you do the GQA adjustment right here, check for smooth_layer.out vs balance_layer.in get GQA factor, then take x_mean and/or w_mean and do the average + expand thing. Then all your scales will be defined correctly going forward

[soyr-redhat]

Now compressing x_mean and w_mean to kv_dim right here. The grid search runs at kv_dim granularity, expanding to hidden_dim inside the loop for W*s on o_proj.

[HDCharles]

1. the scales need to be set during the compute_best_scales step, there's not really a point to doing this unless we can actually calculate the correct output error.

so e.g. when you go to compute the scales for the grid search, you have to check for GQA and then do the average+expand thing when necessary.

2. I'm trying to understand if all the new logic is necessary. It seems like we can just always check for compatible dimensions A.in mod B.out == 0 and treat it like GQA and if a model can't handle that, they can edit the mappings to not include that mapping. So then can we just wait until the point where we generate the scales in compute_best_scales? I don't see why not. But i could be missing something?

[soyr-redhat]

@HDCharles Latest round has been addressed! Would appreciate another look if you had any availability. Thank you!