FlashAttention scaling, multi-GPU training and weight quantization

Modern AI infrastructure relies on hardware-aware attention kernels, distributed multi-GPU parallelism, and low-bit quantization to train and serve massive models efficiently. These three techniques bridge the gap between theoretical model scale and physical GPU limits.

1. FlashAttention-2 and FlashAttention-3 on Hopper (H100) GPUs

While the original FlashAttention solved the memory bandwidth bottleneck, it still left significant raw compute performance unused on modern GPUs. FlashAttention-2 (2023) and FlashAttention-3 (2024) reworked how GPU threads and hardware engines coordinate.

• FlashAttention-2 (Better Work Partitioning).

• Reducing Non-Matmul Operations. GPU Tensor Cores execute matrix multiplications (matmul) roughly 16× faster than standard CUDA cores execute scalar operations like Softmax rescaling. FlashAttention-2 rewrote the online Softmax math to postpone rescaling until the very end of the inner loop, cutting down slow non-matmul instructions.
• Parallelizing Over Sequence Length. The original version parallelized work mainly across batch size and attention heads, leaving GPU cores idle when training with long sequences and small batch sizes. FlashAttention-2 adds parallelism across the sequence length dimension and flips the loop order (outer loop over Q blocks, inner loop over K and V), minimizing slow shared-memory writes.

• FlashAttention-3 (Hopper H100 Hardware Exploitation).

• Asynchronous Data Movement (TMA). NVIDIA's Hopper architecture introduced the Tensor Memory Accelerator (TMA), a dedicated hardware unit that copies data between global HBM and on-chip SRAM without tying up compute threads.
• Warp Specialization (Producer-Consumer Model). FlashAttention-3 splits GPU warps (groups of 32 threads) into dedicated roles. Producer warps use the TMA to prefetch future tiles of K and V into SRAM, while Consumer warps feed current tiles into the Tensor Cores (WGMMA instructions). This overlaps memory transfers with matrix math and Softmax calculations simultaneously.
• FP8 Low-Precision Support. FlashAttention-3 unlocks Hopper's FP8 Tensor Cores, doubling raw throughput over FP16. To prevent precision loss from outlier activations in FP8, it uses block quantization (scaling small tiles independently) and incoherent processing (multiplying Q and K by random orthogonal matrices to smooth out extreme values before quantizing).

2. Multi-GPU Scaling Strategies (3D Parallelism)

When a model has tens or hundreds of billions of parameters, its weights, gradients, and optimizer states exceed the VRAM of a single GPU. Large-scale training combines three distinct parallelization strategies—often called 3D Parallelism.

• Data Parallelism & ZeRO / FSDP. In standard Data Parallelism, every GPU holds a full copy of the model, wasting memory. Modern frameworks use ZeRO (Zero Redundancy Optimizer) or PyTorch FSDP (Fully Sharded Data Parallel), which shards parameters, gradients, and optimizer states across all GPUs, gathering each layer's weights just in time for computation.

• Tensor Parallelism (Intra-Layer). Instead of duplicating a massive feed-forward matrix multiply Y = XW, Tensor Parallelism splits W column-wise across GPUs (W = [W_1, W_2]). Each GPU computes its half of the output simultaneously and exchanges partial results over high-speed NVLink (900 GB/s on H100s).

• Pipeline Parallelism (Inter-Layer). Layers 1–8 sit on GPU 1, Layers 9–16 on GPU 2, and so on. To prevent downstream GPUs from sitting idle while waiting for earlier layers to finish (the "pipeline bubble"), frameworks slice each training batch into smaller micro-batches so all GPUs stay busy in an assembly-line fashion.

3. Quantization (GPTQ, AWQ) and On-the-Fly Dequantization Kernels

During inference (generating tokens one by one), Large Language Models are memory-bandwidth bound, not compute bound. Loading 70 billion 16-bit parameters (140 GB) from VRAM just to generate a single token leaves the Tensor Cores waiting on memory transfers. Quantization compresses weights from 16-bit floats down to 4-bit integers (INT4), cutting memory size and transfer time by nearly 4×.

• GPTQ (Post-Training Quantization via Hessian Calibration).

• Naive rounding of 16-bit weights to 4-bit integers destroys model accuracy. GPTQ feeds a small calibration dataset through the model and uses second-order math (the inverse Hessian matrix) to measure how sensitive each weight is.
• It quantizes weights column-by-column; whenever rounding one weight introduces an error, GPTQ mathematically adjusts the remaining unquantized weights in that row to cancel out the error.

• AWQ (Activation-Aware Weight Quantization).

• Researchers discovered that roughly 0.1% to 1% of weight channels are vastly more important than the rest because they multiply large activation values.
• Instead of keeping those critical weights in 16-bit (which creates messy, slow mixed-precision memory layouts), AWQ scales up those salient weight channels before quantization and scales down the corresponding input activations at runtime, preserving accuracy while keeping 100% of the weights in uniform 4-bit format.

• How Custom CUDA Kernels Dequantize on the Fly (W4A16).

• GPUs cannot directly multiply a 16-bit activation (FP16) by a 4-bit integer weight (INT4). Instead, specialized inference kernels (such as Marlin or ExLlama) perform W4A16 (4-bit Weights, 16-bit Activations) execution:
• 1. Compressed HBM Read: The kernel loads packed 4-bit weights (eight INT4 values packed inside a single 32-bit integer) from slow global VRAM, slashing memory bandwidth usage by 4×.
• 2. Register-Level Unpacking: Inside the GPU's ultra-fast registers, the CUDA kernel uses fast bitwise shifts (>>) and bitmasks (&) to unpack the INT4 integers and applies stored scale/zero-point factors to convert them back into FP16 values on the fly.
• 3. Immediate Tensor Core Math: The freshly unpacked FP16 weights are fed straight into the Tensor Cores alongside the FP16 activations without ever writing the expanded 16-bit weights back to VRAM.