AWQ Quantization Support

Status: Implemented (P0) Issue: #485 (A-001)

AWQ (Activation-aware Weight Quantization) is a 4-bit quantization method that achieves near-float performance by using activation-aware calibration to determine optimal per-channel scaling factors.

Who this is for: engineers loading or serving AWQ-quantized safetensors with fak, or exporting their own AWQ checkpoints. Prerequisites: familiarity with 4-bit quantization basics (codes, scales, zero-points) and Go for the loader snippets. By the end you’ll know fak’s on-disk AWQ layout and dequant formula, how to call model.LoadAWQ, and how to produce a checkpoint with AutoAWQ.

Overview

AWQ reduces model memory footprint to ~0.5625 bytes per parameter (4-bit weights + per-channel scales) compared to:

AWQ achieves this while maintaining >99% of FP32 accuracy through activation-aware scale calibration.

Format Specification

Data Layout

Dequantization Formula

weight = scale[o] × (code - 8)

Where code is the unpacked 4-bit value (0-15) and 8 is the symmetric zero-point.

Usage

Loading AWQ Models

import "github.com/anthony-chaudhary/fak/internal/model"

// Load from directory containing model.safetensors with AWQ weights
m, err := model.LoadAWQ("/path/to/awq/model")
if err != nil {
    log.Fatal(err)
}

// Check AWQ tensors loaded
fmt.Printf("Loaded %d AWQ tensors\n", m.AWQCount())

AWQ Tensor Format

AWQ quantized safetensors use the following naming convention:

For example, for a QKV projection:

model.layers.0.self_attn.q_proj.weight      # 4-bit packed weights
model.layers.0.self_attn.q_proj.weight_scale # scales

Integration with Forward Pass

The AWQ kernel provides:

// Matrix-vector multiplication: y = A @ x
y := awqMatRows(awqTensor, x)

// Batched matmul: Y = A @ X^T (P tokens)
Y := awqGemm(awqTensor, X, P)

Creating AWQ Checkpoints

Using AutoAWQ (Python)

from autoawq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_path = "meta-llama/Llama-3.1-8B"
quant_path = "./llama-3.1-8b-awq"

quantizer = AutoAWQForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)

quantizer.quantize(tokenizer, quant_config={
    "zero_point": True,
    "q_group_size": 128,
    "n_sample_calib": 32,
})

quantizer.save_quantized(quant_path)
Model Group Size Calibration Samples
Llama 2/3 128 32
Qwen2/2.5 128 32
Mistral 128 64

Performance

Memory Savings

| Model | FP32 | AWQ | Reduction | |——-|——|—–|————| | Llama-3.1-8B | 16 GB | 4.5 GB | 3.6× | | Llama-3.1-70B | 140 GB | 40 GB | 3.5× | | Qwen2.5-7B | 14 GB | 4 GB | 3.5× |

Accuracy

AWQ typically achieves >99% of FP32 accuracy on standard benchmarks:

Throughput

Decode speed depends on backend:

Implementation Details

CPU Kernels

CUDA Kernels

The CUDA implementation computes the matmul directly on packed 4-bit weights without full dequantization, achieving near-Q8 throughput with ~3.5× memory savings.

Testing

Oracle tests verify:

Run tests:

go test -v -run TestAWQ ./internal/model/...

Limitations

  1. CUDA requires rebuild — Must compile with -tags cuda (uses cgo)
  2. Requires even input dimensions — Padded by AWQ export
  3. No zero-point tensors — Assumes symmetric quantization
  4. Safetensors only — Pytorch bin format not yet supported

Future Work

  1. AVX2/AVX-512 assembly kernels — For faster CPU dequantization
  2. CUDA graph integration — Capture AWQ ops in decode graph
  3. Mixtral AWQ — MoE models with AWQ quantization
  4. Dynamic AWQ — Runtime quantization without pre-export

References