[Bugfix] Enable attn quantization of Llama-4 by correctly permuting scales for rope (int8, fp8)

[eldarkurtic]

Llama-4 weights of q/k_proj are permuted during model loading to prepare the model for interleaved/gpt-neox rope.
The same permutation needs to be applied on quantization weight scales as well.

Purpose 1

So far, all of quantized Llama-4 models used to skip attention quantization as the model's accuracy would collapse. This was mistakenly interpreted as high sensitivity of the attention layers under quantization. Luckily for quantization, this is not the case and the accuracy collapse was caused by the mismatch in weights and their scales after weight-permutations for rope.

After this PR, we can safely quantize attention layers and maintain competitive accuracy relative to the unquantized baseline.

GSM8k evaluation command:

lm_eval --model vllm --model_args pretrained=MDL,add_bos_token=True,tensor_parallel_size=4,max_model_len=10000 --tasks gsm8k --num_fewshot 5 --batch_size auto

Results:

meta-llama/Llama-4-Scout-17B-16E-Instruct                                                   = 90.67
RedHatAI/Llama-4-Scout-17B-16E-Instruct-FP8-dynamic (w/o attention quant)                   = 90.90
ekurtic/Llama-4-Scout-17B-16E-Instruct-FP8-dynamic_wAttn (w/ attention quant, w/o this PR)  =  5.38
ekurtic/Llama-4-Scout-17B-16E-Instruct-FP8-dynamic_wAttn (w/ attention quant, w/ this PR)   = 91.13

Purpose 2

Enable quantization of meta-llama/Llama-Guard-4-12B models. Not quantizing attention layers in Llama-4 MoEs wasn't a big penalty as those weights usually amount to 5-10% of the total weights compared to experts. However, Llama-Guard-4 models are built by replacing MoE part of Llama-4 with a standard dense layer, resulting in a standard decoder model where attention layers amount to a noticeable fraction of total weights (>10%) and therefore we would prefer to quantize them. This PR enables accurately quantizing all Linear layers (self-attn and MLPs) for Llama-Guard models, just as it is done for all other standard decoder-only architectures.

Note: for simplicity, this PR is a bugfix for INT8 and FP8 per-channel quant schemes. Patch for INT4 will come in a follow-up PR, as it requires slightly more complicated logic due to packed weights.

[gemini-code-assist]

Code Review

This pull request effectively addresses the bug where attention quantization of Llama-4 models led to accuracy collapse due to incorrect permutation of quantization scales. The changes correctly extend the weight permutation logic to handle scales for INT8 and FP8 quantization schemes, ensuring competitive accuracy relative to the unquantized baseline. This is a crucial fix for enabling quantized attention layers in Llama-4 and Llama-Guard-4 models.

[mergify]

Hi @eldarkurtic, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[mgoin]

Why is this needed, because CT transposes scales? Could we just always use the w.shape to decide?

[eldarkurtic]

Notice that the original implementation was: attn_out = self.config.hidden_size and there is a comment above which says Do not rely on w's shape, as it may be in another layout. I didn't want to break this by doing attn_out = w.shape[-1].

[eldarkurtic]

If you think this is safe to ignore, and always use attn_out = w.shape[-1], let me know

[mgoin]

I would like to avoid CT specific logic here. Why don't we apply to all weight scales? Is it because of concern over packed weights? I think they wouldn't work anyway

[eldarkurtic]

I added CT specific logic here for safety reasons. I don't know how other frameworks are shipping weight-scales