Kernels & Compilers advanced 9 min read 7 flashcards

Quantised GEMM Kernels

Quantised GEMM kernels replace 16-bit or 32-bit matrix multiplications with 8-bit or 4-bit integer arithmetic, cutting memory bandwidth and compute cost while preserving model accuracy through careful scaling and outlier handling.

A single forward pass of Llama-3 70B in FP16 moves roughly 140 GB of weight data through GPU memory each token. At 2 TB/s of HBM bandwidth, that is 70 ms/token before a single multiply-add has been computed. Switching to INT8 halves the transfer; INT4 quarters it. Quantised GEMM kernels are the mechanism that actually realises those savings at the hardware level, and getting them right is harder than changing a dtype flag.

What a GEMM kernel really does

A general matrix multiplication (GEMM) computes C = A × B, where the operands live in GPU global memory (HBM) and the result is written back. The kernel's job is to:

  1. Tile the matrices into blocks that fit in on-chip SRAM (shared memory on CUDA devices).
  2. Load a tile of A and a tile of B cooperatively across the thread block.
  3. Accumulate partial dot products in registers.
  4. Write the result tile back to global memory.

In FP16 GEMM, both storage and accumulation can happen at the same precision. In quantised GEMM the workflow splits:

Stage Typical dtype
Weight storage (on disk / HBM) INT4 or INT8
Dequantisation FP16 or BF16
Matrix multiply (tensor core accumulator) FP16/BF16 or INT32
Output / residual FP16 or BF16

The key insight is that GPU tensor cores on Ampere/Hopper can execute INT8 × INT8 → INT32 natively at roughly 2x the throughput of FP16 (same number of tensor core cycles, but twice as many INT8 elements fit per cycle). INT4 support varies by generation: Ada and Hopper expose FP8 tensor cores, which is often more practical than INT4 for inference.

The dequantisation bottleneck

Naive quantisation stores every weight as INT8 and multiplies by a single global scale factor before the GEMM. This fails in practice for two reasons.

Outlier features. From roughly 6.7B parameters upward, transformer hidden states develop systematic high-magnitude channels (outlier features). If you quantise activations with a global scale chosen to cover the outlier, the other 99% of values are crushed into a handful of quantisation bins, destroying accuracy. LLM.int8() (Dettmers et al., 2022) handles this by decomposing each matrix multiplication: the outlier columns are isolated and kept in FP16, and the remaining 99.9%+ of values are multiplied in INT8. The two partial products are summed in FP16.

Per-channel vs. per-tensor scaling. A single scale per tensor wastes range. Per-channel (or per-group) scales let each output channel choose its own dynamic range. Most production kernels use group quantisation with a group size of 64 or 128 weights sharing one scale and zero-point:

# Conceptual layout: weight shape [K, N], group_size = 128
# scales shape: [K // 128, N]
# actual weight (decompressed):
w_fp16[k, n] = (w_int4[k, n] - zero_point[k // 128, n]) * scale[k // 128, n]

The dequantise-then-multiply path is how AWQ (Lin et al., 2023) and GPTQ (Frantar et al., 2022) both work in practice: weights stay INT4 in HBM, the kernel unpacks and scales each group just before the tensor core computation, and accumulation happens in FP16 or BF16. Bandwidth is determined by INT4 traffic; compute is determined by FP16 throughput.

SmoothQuant and the W8A8 path

The approaches above are weight-only quantisation: activations stay in FP16. That still halves the weight bandwidth but does nothing to reduce the multiply cost when activations dominate. True INT8 × INT8 GEMMs (W8A8) need both weights and activations to be quantisable without catastrophic outlier damage.

SmoothQuant (Xiao et al., 2023) solves this with a mathematically equivalent transformation. If activation channel c has high variance, multiply the weights in column c by a scale factor s_c and divide the activations by the same factor. The product is unchanged, but the difficulty has been shifted from activations (which are dynamic and hard to calibrate) to weights (which are fixed and easy to absorb offline):

Y = (X / s) × (W × s)   # equivalent to Y = X × W

After smoothing, both X and W have roughly similar per-channel magnitudes, and a standard INT8 quantiser can cover them without wasted range. With W8A8, the full GEMM runs on INT8 tensor cores, giving up to 1.56x speedup and 2x memory reduction versus FP16 on a range of transformer models.

Kernel-level implementation details

Once the mathematical scheme is chosen, the kernel author faces several engineering constraints.

Vectorised loads. INT4 weights are packed two-per-byte (or four-per-byte for INT2). Loading them efficiently requires 128-bit (or 256-bit) vectorised loads using int4 or uint4 CUDA intrinsics, then bit-shifting and masking to extract individual nibbles before applying scale.

Register pressure. Each thread holds partial sums across multiple accumulator tiles. Unpacking INT4 to FP16 before the tensor core operation inflates register count; spilling to local memory destroys throughput. CUTLASS and libraries like marlin carefully orchestrate the pipeline so that unpacking and the GEMM overlap in the instruction stream.

Asymmetric vs. symmetric quantisation. Symmetric quantisation sets zero-point to zero, simplifying the kernel (no subtraction). Asymmetric quantisation uses a non-zero zero-point for better coverage of rectified activations (e.g., post-ReLU or post-SiLU), but the zero-point subtraction must be fused into the scale application to avoid extra passes.

Batched vs. single-token inference. At batch size 1 (autoregressive decoding), the GEMM degenerates to a GEMV (matrix-vector product). The bottleneck shifts entirely to memory bandwidth; compute utilisation is near zero. Quantisation to INT4 is most valuable here because HBM bandwidth savings map directly to latency reduction. For large batches (prefill or offline throughput), compute becomes relevant again and W8A8 INT8 schemes are competitive.

When it falls down

Accuracy degradation at low bit-widths. INT4 per-tensor quantisation of activations is often too coarse for large models. Models that rely on extreme outlier channels (common in GPT-style architectures) need specialised handling. Below 4 bits, standard round-to-nearest quantisation breaks most models without extensive calibration or quantisation-aware training.

Kernel fragmentation. The ecosystem has dozens of kernel implementations: bitsandbytes (used by LLM.int8()), marlin, awq-gemm, CUTLASS INT8 GEMMs, TensorRT's QDQ graphs, and vLLM's quantised paged attention. Each targets a different combination of GPU generation, batch size regime, and quantisation scheme. A kernel that is fastest on Ampere for INT8 may be slower than FP16 on Hopper for the same workload because Hopper's FP8 tensor cores offer a different trade-off.

Calibration sensitivity. Per-channel and per-group scales are estimated from a small calibration dataset. If the calibration distribution does not cover activation outliers that appear in deployment (e.g., very long documents, unusual token distributions), the scales are miscalibrated and accuracy degrades silently.

Hardware compatibility. INT8 tensor cores require SM 7.5 (Turing) or later. INT4 tensor cores are SM 8.9+ (Ada Lovelace) for practical use. Quantised GEMM kernels written for Hopper (FP8 native, NV TMA async copies) do not run on older hardware at all; the code paths are completely separate.

Memory layout constraints. Most quantised kernels require weights in a specific tiled layout (row-major, column-major, or interleaved tile format) that differs from the layout produced by standard training. A runtime repack is needed on first load, which adds latency that matters in serverless deployments with cold starts.

Further reading

Check yourself

7 flashcards for this concept

Click a card to reveal the answer.

Drill the whole track