Quantization Hardware: Tensor Cores & Custom Kernels

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. NVIDIA Tensor Core Architecture
5. INT8 Tensor Cores (Turing/Ampere)
6. FP8 Tensor Cores (Hopper)
7. 2:4 Structured Sparsity (Ampere)
8. CPU Quantized Inference (SIMD)
9. Custom Kernels for Mixed Precision
10. Apple & Mobile Accelerators
11. Common Pitfalls
12. Summary
13. Exercises

Learning Objectives

1. Explain how tensor cores execute quantized matrix multiply.
2. Describe the data layout requirements for INT8/FP8 tensor cores.
3. Explain 2:4 structured sparsity hardware support.
4. Describe CPU SIMD optimizations for quantized models (AVX2, ARM NEON).
5. Choose quantization format based on target hardware.

Notation

• MMA — matrix multiply-accumulate (tensor core operation)
• TOPS — tera operations per second
• SIMD — single instruction, multiple data

Core Intuition

Quantization only provides speedup if the hardware can exploit the reduced precision. A GPU that can only do FP16 matmul sees no benefit from INT4 weights stored in memory (still needs to dequantize). Real speedup requires: (1) native low-precision compute units (tensor cores for INT8/FP8), or (2) SIMD parallelism where more operations fit per instruction (CPU vectorization). Hardware-software co-design determines which quantization methods are practical.

NVIDIA Tensor Core Architecture

Tensor cores perform small matrix multiply-accumulate operations in a single clock cycle:

where $\mathbf{A}, \mathbf{B}$ are in low precision and $\mathbf{C}, \mathbf{D}$ are in higher precision (accumulator).

Supported configurations (H100):

Key observation: INT8 and FP8 deliver exactly 2x the throughput of FP16/BF16.

INT8 Tensor Cores (Turing/Ampere)

Operation: $D_{\text{INT32}} = A_{\text{INT8}} \times B_{\text{INT8}} + C_{\text{INT32}}$.

Tile size (A100): $16 \times 8 \times 32$ (M=16, N=8, K=32 per tensor core MMA).

Workflow for W8A8:

1. Quantize weights offline (INT8 per-channel).
2. Quantize activations online (INT8 per-token).
3. INT8 matmul on tensor cores → INT32 accumulator.
4. Rescale INT32 output by $(s_W \times s_A)$ → FP16/BF16 output.

Data layout requirement: Column-major for B, row-major for A. Incorrect layout → falls back to CUDA cores (10x slower).

Frameworks: CUTLASS, cuBLAS INT8 GEMM, TensorRT-LLM.

FP8 Tensor Cores (Hopper)

Advantage over INT8:

• No explicit quantize/dequantize step — just scale and cast.
• Same throughput (1979 TOPS on H100).
• Simpler kernel implementation.
• Native support for mixed E4M3/E5M2 inputs.

NVIDIA Transformer Engine: Automatic FP8 management:

1. Tracks per-tensor scale history.
2. Applies delayed scaling.
3. Handles overflow detection and scale correction.
4. Transparent to the user (just set fp8_autocast).

2:4 Structured Sparsity (Ampere)

Hardware support (A100+): Specialized sparse tensor core format.

Pattern: Exactly 2 zeros per group of 4 elements (50% sparsity):

$\binom{4}{2} = 6$ possible patterns per group.

Storage: Compressed format stores only the 2 non-zero values + a 2-bit index encoding which positions are non-zero.

Speedup: 2x over dense at same precision (tensor core skips zero multiplications).

Combined with quantization: INT8 + 2:4 sparsity = 4x effective speedup over FP16 dense.

CPU Quantized Inference (SIMD)

AVX2 (x86-64): 256-bit registers → 32 INT8 operations per instruction.

AVX-512 VNNI: Dedicated dot-product instruction for INT8:

• vpdpbusd: Multiply 64 pairs of INT8 values, accumulate to INT32.
• 4x throughput over scalar.

ARM NEON (Apple M-series, mobile): 128-bit registers → 16 INT8 per instruction.

• sdot/udot: Dot product of 4xINT8 → INT32.

llama.cpp optimization: Hand-tuned SIMD kernels for each quantization type:

• Q4_K: Dequantize 32 weights at once, multiply with FP32 activations.
• Q8_0: Direct INT8×INT8 with VNNI.
• Achieves 60-90% of theoretical peak throughput.

Custom Kernels for Mixed Precision

W4A16 dequantization kernel (vLLM/Marlin):

1. Load INT4 weights (packed: 2 per byte).
2. Dequantize to FP16 in registers.
3. FP16 tensor core matmul.
4. Repeat for next tile.

Marlin kernel (optimized W4A16): Achieves near-optimal throughput by:

• Overlapping dequantization with matmul (software pipelining).
• Optimal shared memory layout.
• 3.6x speedup over naive W4A16 on A100.

Key insight: Even without native INT4 tensor cores, W4A16 provides speedup through reduced memory bandwidth (weight loading is the bottleneck at small batch sizes).

Apple & Mobile Accelerators

Apple Neural Engine (ANE):

• Native INT8 support with per-channel quantization.
• CoreML framework handles quantization automatically.
• 15.8 TOPS (M1), 38 TOPS (M4).

Qualcomm Hexagon DSP:

• INT8 and INT4 support.
• Used for on-device LLM inference (Snapdragon 8 Gen 3).

Google TPU:

• INT8 matrix units.
• BF16 native support (Google invented BF16).

Common Pitfalls

Pitfall 1. Assuming INT4 weights provide automatic 2x speedup on GPU. Without native INT4 tensor cores (available on Blackwell, not Hopper), W4A16 relies on dequantization + FP16 matmul. Speedup is from memory bandwidth, not compute.

Pitfall 2. Incorrect data layout for tensor cores. If weight matrix isn't in the required column-major/interleaved format, CUDA falls back to non-tensor-core path. Always verify with profiler.

Pitfall 3. Ignoring CPU quantization for small models. For models under 7B, CPU inference with optimized SIMD kernels (llama.cpp) can be faster than GPU with kernel launch overhead.

Summary

• GPU tensor cores: INT8/FP8 at 2x throughput over FP16 (A100/H100).
• 2:4 sparsity: Additional 2x with structured pruning on Ampere+.
• CPU SIMD: AVX2/VNNI/NEON for efficient quantized inference without GPU.
• Custom kernels (Marlin, CUTLASS) bridge the gap for mixed-precision formats.
• Hardware determines viable formats: Match quantization to target hardware.
• W4A16: Memory-bound speedup (bandwidth). W8A8: Compute-bound speedup (tensor cores).

Exercises

Exercise 1. Compute the peak INT8 throughput of an A100 in TOPS. Compare to FP16 TFLOPS. What's the theoretical speedup for a compute-bound workload?

Exercise 2. For 2:4 sparsity on a $4096 \times 4096$ matrix: compute the compressed size (non-zeros + indices) vs dense storage.

Exercise 3. Design an AVX-512 loop for INT8 matmul of a 1024x1024 matrix. Estimate throughput on a modern CPU (assume 2GHz, 2 VNNI units).

Exercise 4. For the Marlin W4A16 kernel: estimate the arithmetic intensity (FLOPs/byte) and determine if it's compute-bound or memory-bound on an A100.

Exercise 5. An edge device has: 8 TOPS INT8, 2 TFLOPS FP16, 50 GB/s bandwidth. For a 3B model: should you use W8A8, W4A16, or W4A8? Justify with roofline analysis.