Diving into the GPUs — Fusing, Threading, and Mixing

1. GPU Architecture Primer

1.1 Compute Hierarchy

An NVIDIA GPU is organized as a hierarchical array of compute units. The fundamental building block is the Streaming Multiprocessor (SM), each of which contains multiple streaming processors (cores).

For the NVIDIA H100 SXM:

Each core can execute multiple threads concurrently. The compute hierarchy from bottom to top is:

1.2 Memory Hierarchy

The memory system is equally hierarchical, with a fundamental trade-off: smaller memories are faster but private; larger memories are slower but shared.

The performance optimization objective is:

1.3 Kernel Execution Model

A kernel is a function that runs on the GPU. It is written in CUDA (C/C++ extension) or Triton (Python-based) and compiled to PTX (Parallel Thread Execution) — NVIDIA’s low-level virtual ISA.

The execution requires two components:

Host Code (CPU side): Allocates device memory, transfers data, launches kernels.

// Allocate on GPU
cudaMalloc(&d_A, size);
// Copy host → device
cudaMemcpy(d_A, h_A, size, cudaMemcpyHostToDevice);
// Launch kernel with grid/block dimensions
VecAdd<<<blocksPerGrid, threadsPerBlock>>>(d_A, d_B, d_C, N);
// Copy device → host
cudaMemcpy(h_C, d_C, size, cudaMemcpyDeviceToHost);

Device Code (GPU side): The kernel function itself, executed by each thread.

__global__ void VecAdd(float* A, float* B, float* C, int N) {
    int i = blockDim.x * blockIdx.x + threadIdx.x;
    if (i < N)
        C[i] = A[i] + B[i];
}

The global thread index for a 1D grid is computed as:

The grid dimensions are chosen to cover all $N$ elements:

1.4 Scheduling Constraints

Key hardware constraints govern efficient scheduling:

Occupancy — the ratio of active warps to the maximum possible — determines how well the SM hides memory latency:

Higher occupancy provides more warps to switch between during memory stalls, hiding latency through latency hiding via warp scheduling.

2. Improving Performance with Kernels

2.1 Toolchain Spectrum

There exists a spectrum of tools for writing GPU kernels, ordered by ease-of-use vs. control:

2.2 Example: Exponential Linear Unit (ELU)

The ELU activation function is defined as:

where $\alpha$is typically set to$1.0$.

PyTorch with @torch.compile:

@torch.compile
def elu(x, alpha=1.0):
    return torch.where(x < 0, alpha * (torch.exp(x) - 1), x)

The @torch.compile decorator invokes TorchInductor, which:

1. Traces the PyTorch operations into an intermediate representation (FX graph)
2. Lowers the graph to Triton kernels
3. Compiles and caches the optimized kernels

The generated Triton kernel can be inspected by setting TORCH_LOGS="output_code":

@triton.jit
def elu_kernel(input_ptr, output_ptr, num_elements, BLOCK_SIZE: tl.constexpr):
    block_start = tl.program_id(0) * BLOCK_SIZE
    block_indices = block_start + tl.arange(0, BLOCK_SIZE)[:]
    valid_mask = block_indices < num_elements
    input_values = tl.load(input_ptr + block_indices, valid_mask)
    zero_value = 0.0
    negative_mask = input_values < zero_value
    exp_values = tl.math.exp(input_values)
    one_value = 1.0
    shifted_exp_values = exp_values - one_value
    output_values = tl.where(negative_mask, shifted_exp_values, input_values)
    tl.store(output_ptr + block_indices, output_values, valid_mask)

Key Triton concepts:

• tl.program_id(0): Returns the unique block ID in dimension 0 (analogous to blockIdx.x in CUDA)
• BLOCK_SIZE: Compile-time constant defining how many elements each block processes
• tl.arange(0, BLOCK_SIZE): Creates a vector of consecutive indices within a block
• valid_mask: Bounds-checks to prevent out-of-range memory access
• tl.load / tl.store: Masked memory operations

Triton vs. CUDA control granularity:

3. CUDA Optimization Techniques

3.1 Memory Coalescing

Definition: Memory coalescing is the hardware mechanism by which a GPU combines multiple memory access requests from threads within a single warp into a minimal number of memory transactions, exploiting the burst transfer behavior of DRAM.

DRAM burst mechanism: When any single address $M$in global memory (HBM/DRAM) is accessed, the DRAM chip reads a contiguous segment of$B_{\text{burst}}$bytes (typically 32 or 128 bytes) around$M$ in a single operation. Coalescing ensures threads access addresses within the same burst segment.

Coalescing condition: For a warp of 32 threads accessing addresses $a_0, a_1, \ldots, a_{31}$, memory is perfectly coalesced when:

That is, consecutive threads access consecutive memory locations.

Naive Matrix Multiplication (Uncoalesced)

Consider computing $\mathbf{C} = \mathbf{A} \cdot \mathbf{B}$where$\mathbf{A} \in \mathbb{R}^{M \times K}$, $\mathbf{B} \in \mathbb{R}^{K \times N}$, $\mathbf{C} \in \mathbb{R}^{M \times N}$.

Row-major storage convention: An element at row $r$, column $c$of a matrix with$N_{\text{cols}}$ columns is stored at linear address:

Naive 2D kernel:

__global__ void matmul_naive(int M, int N, int K, 
                              const float *A, const float *B, float *C) {
    const uint x = blockIdx.x * blockDim.x + threadIdx.x;  // row
    const uint y = blockIdx.y * blockDim.y + threadIdx.y;  // column
    if (x < M && y < N) {
        float tmp = 0.0;
        for (int i = 0; i < K; ++i) {
            tmp += A[x * K + i] * B[i * N + y];
        }
        C[x * N + y] = tmp;
    }
}

Problem analysis: Two threads in the same warp with consecutive threadIdx.x values (e.g., thread $(0,0)$and thread$(1,0)$) have the same $y$but different$x$. At iteration $i = 0$:

• Thread $(0, 0)$reads$A[0 \cdot K + 0] = A[0]$
• Thread $(1, 0)$reads$A[1 \cdot K + 0] = A[K]$

These addresses are $K$elements apart in memory — not consecutive. The accesses to$\mathbf{A}$are uncoalesced, resulting in$\sim 32$ separate memory transactions instead of 1.

Fix — 1D block with recomputed indices:

const int x = blockIdx.x * BLOCKSIZE + (threadIdx.x / BLOCKSIZE);  // row
const int y = blockIdx.y * BLOCKSIZE + (threadIdx.x % BLOCKSIZE);  // column

Now consecutive threadIdx.x values produce the same $x$(row) but different$y$(column). At iteration$i = 0$:

• Thread 0 reads $A[x \cdot K + 0]$and$B[0 \cdot N + y_0]$
• Thread 1 reads $A[x \cdot K + 0]$(same!) and$B[0 \cdot N + y_0 + 1]$

The $\mathbf{B}$accesses are now consecutive in memory (coalesced), and the$\mathbf{A}$accesses are identical (broadcast). The result is a $\sim 10\times$ improvement in both memory throughput and execution time.

3.2 Tiling (Shared Memory Optimization)

Motivation: Even with coalesced access, global memory bandwidth is limited. For matrix multiplication, each element of $\mathbf{A}$and$\mathbf{B}$ is loaded multiple times by different threads. Without optimization, the total number of global memory loads is:

(Each of the $M \times N$output elements requires$K$loads from$\mathbf{A}$and$K$loads from$\mathbf{B}$.)

Tiling principle: Partition the computation into tiles of size $T_M \times T_K$(from$\mathbf{A}$) and $T_K \times T_N$(from$\mathbf{B}$). A block of threads cooperatively loads one tile of each matrix into shared memory (SRAM), then all threads in the block compute using the fast shared memory.

The output tile $\mathbf{C}_{\text{tile}} \in \mathbb{R}^{T_M \times T_N}$is accumulated over$\lceil K / T_K \rceil$ iterations:

For a square tile of size $T$:

for (int tileIdx = 0; tileIdx < K; tileIdx += TILE_SIZE) {
    // Cooperative load: each thread loads one element of A and B
    sharedA[localRow * TILE_SIZE + localCol] = A[localRow * K + localCol];
    sharedB[localRow * TILE_SIZE + localCol] = B[localRow * N + localCol];
    __syncthreads();  // Barrier: ensure all loads complete
    
    // Compute partial dot product from shared memory
    for (int i = 0; i < TILE_SIZE; ++i) {
        sum += sharedA[localRow * TILE_SIZE + i] 
             * sharedB[i * TILE_SIZE + localCol];
    }
    __syncthreads();  // Barrier before next tile load
    
    A += TILE_SIZE;       // Advance tile in A (across columns)
    B += TILE_SIZE * N;   // Advance tile in B (down rows)
}

Memory access reduction: With tiling, each element of $\mathbf{A}$and$\mathbf{B}$is loaded from global memory only once per tile, then reused$T$ times from shared memory:

The reuse factor is $T$, giving a proportional reduction in global memory traffic.

__syncthreads() is a block-level barrier: all threads in the block must reach this point before any can proceed. Two barriers are required per tile iteration:

1. After loading data into shared memory (ensure all data is available before computation)
2. After computation (ensure all threads finish before overwriting shared memory with the next tile)

Shared memory requirement:

For $T_M = T_K = T_N = T = 32$ with FP32:

This easily fits within the 228 KB per SM on H100.

3.3 Thread Coarsening

Problem identified via profiling: After tiling, warp stall analysis reveals that warps spend significant cycles in the stalled_mio_throttle state — stalled waiting for the Memory Input/Output (MIO) pipeline to service shared memory requests.

The root cause is that each thread computes a single output element, requiring many shared memory loads per thread. The shared memory access pipeline becomes a bottleneck.

Thread coarsening merges $C_f$threads into a single coarsened thread, where each coarsened thread computes$C_f$ output elements instead of 1.

Before coarsening: Each thread computes 1 element, requiring $2T_K$ shared memory loads per tile iteration:

(One row of $\mathbf{A}_{\text{shared}}$, one column of $\mathbf{B}_{\text{shared}}$.)

After coarsening by factor $C_f$: Each thread computes $C_f$elements in the same row, sharing the same row of$\mathbf{A}_{\text{shared}}$:

Without coarsening, $C_f$ separate threads would have loaded:

The savings ratio is:

For $C_f = 8$: reduction factor $= 16/9 \approx 1.78\times$ fewer shared memory accesses.

3.4 Minimizing Control Divergence

GPUs execute in the SIMD (Single Instruction, Multiple Data) model at the warp level: all 32 threads in a warp execute the same instruction simultaneously on different data.

Control divergence occurs when threads within a warp encounter a conditional branch and take different execution paths:

if (condition) {
    // Path A — executed by some threads
} else {
    // Path B — executed by remaining threads
}

When divergence occurs, the hardware serializes execution:

1. Threads taking Path A execute while Path B threads are masked (idle)
2. Then Path B threads execute while Path A threads are masked

The effective warp throughput drops proportionally:

For a simple if-else, $N_{\text{paths}} = 2$, yielding 50% throughput.

Mitigation strategies:

4. Fused Kernels

4.1 The Kernel Launch Overhead Problem

In standard (unfused) execution, each operation is a separate kernel launch. For a sequence of $n$point-wise operations$f_1, f_2, \ldots, f_n$on a tensor$\mathbf{x}$:

Each kernel launch requires:

1. Write $\mathbf{x}_i$ from SM registers/shared memory → global memory (HBM)
2. Kernel launch overhead (CPU→GPU scheduling, ~5–10 μs)
3. Read $\mathbf{x}_i$ from global memory → SM for the next kernel

The total HBM traffic for $n$ unfused kernels:

(Each intermediate result is written once and read once from HBM.)

4.2 Kernel Fusion

Definition: Kernel fusion combines multiple operations into a single kernel that executes all computations without materializing intermediate results in global memory.

The fused HBM traffic:

The bandwidth savings are:

For LayerNorm, which involves ~5 point-wise operations (subtract mean, square, average, reciprocal square root, scale+shift), the savings factor is $\sim 4\times$.

4.3 Applicability in Transformers

Fusion is most beneficial for memory-bound operations — those where HBM bandwidth, not compute throughput, is the bottleneck. The arithmetic intensity (AI) determines this:

Fusion provides the greatest speedup for memory-bound operations since reducing HBM traffic directly reduces the bottleneck.

5. FlashAttention

5.1 Standard Attention and Its Memory Bottleneck

The standard scaled dot-product attention for a single head is:

where $\mathbf{Q}, \mathbf{K}, \mathbf{V} \in \mathbb{R}^{N \times d_k}$, $N$is the sequence length, and$d_k$ is the head dimension.

The naive computation proceeds in three steps:

1. Compute score matrix: $\mathbf{S} = \frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d_k}} \in \mathbb{R}^{N \times N}$
2. Compute attention weights: $\mathbf{P} = \text{softmax}(\mathbf{S}) \in \mathbb{R}^{N \times N}$
3. Compute output: $\mathbf{O} = \mathbf{P}\mathbf{V} \in \mathbb{R}^{N \times d_k}$

Memory problem: Both $\mathbf{S}$and$\mathbf{P}$ must be materialized in HBM. Their size is:

For $N = 4096$, $d_k = 128$, in BF16:

With $n_h = 96$ heads (as in a large model):

This is a significant fraction of the 80 GB HBM on an H100, and the HBM read/write traffic for these matrices becomes the dominant bottleneck.

HBM traffic for naive attention: Each of the three steps requires reading inputs from HBM and writing outputs back:

The dominant terms are all $O(N^2)$, making naive attention memory-bandwidth-bound for typical $N \gg d_k$.

5.2 FlashAttention Algorithm

Core idea: Compute the output $\mathbf{O}$by processing$\mathbf{Q}$, $\mathbf{K}$, $\mathbf{V}$in tiles that fit in SRAM, never materializing the full$N \times N$matrices$\mathbf{S}$or$\mathbf{P}$ in HBM.

The key mathematical challenge is that softmax requires global normalization across the entire row:

FlashAttention uses the online softmax algorithm (Milakov & Gimelshein, 2018) to compute softmax incrementally over tiles. The algorithm maintains running statistics $m_i$(row-wise maximum) and$\ell_i$ (row-wise sum of exponentials) that are updated as each new tile is processed.

Tiled computation: Partition $\mathbf{K}$and$\mathbf{V}$into blocks of$B_c$rows (columns of the attention matrix), and$\mathbf{Q}$into blocks of$B_r$ rows:

where $T_c = \lceil N / B_c \rceil$.

For each query block $\mathbf{Q}_i$($B_r$rows), iterate over KV blocks$j = 1, \ldots, T_c$:

Step 1: Compute the local score tile in SRAM:

Step 2: Compute local row-wise maximum and update running maximum:

Step 3: Compute local exponentials and update running sum:

Step 4: Update the output accumulator:

After all $T_c$iterations,$\mathbf{O}_i^{\text{final}}$ equals the exact attention output — no approximation is involved.

Tile sizes are chosen to maximize SRAM utilization:

5.3 Complexity Analysis

Since $M_{\text{SRAM}} \gg d_k$typically, the HBM access is reduced by a factor of approximately$M_{\text{SRAM}} / d_k$. For H100 with $M_{\text{SRAM}} \approx 228$KB and$d_k = 128$in BF16 ($= 256$ bytes per row):

5.4 Impact and Significance

FlashAttention resolves multiple bottlenecks simultaneously:

1. Memory reduction: The $O(N^2)$memory for the attention matrix is eliminated, reducing to$O(N)$ auxiliary storage.

2. Wall-clock speedup: By reducing HBM traffic — the true bottleneck for attention — FlashAttention achieves 2–4× speedup over standard attention implementations despite performing the same number of FLOPs.

3. Enabling longer sequences: Without the $N^2$ memory overhead, much longer sequences become feasible. This is why earlier linear/subquadratic attention approximations have been largely abandoned in favor of FlashAttention — exact attention is now fast enough.

5.5 FlashAttention-2 and FlashAttention-3

6. Mixed Precision Training

6.1 Floating-Point Number Representation

A floating-point number $x$ in IEEE 754 format is represented as:

where:

• $s \in \{0, 1\}$: sign bit
• $E$: stored exponent value (unsigned integer)
• $\text{bias} = 2^{e-1} - 1$: exponent bias ($e$ = number of exponent bits)
• $p$: number of mantissa bits
• $b_i$: individual mantissa bits

The three components control distinct properties:

6.2 Format Comparison

Machine epsilon $\epsilon$is defined as the smallest$\epsilon > 0$such that$\text{fl}(1 + \epsilon) \neq \text{fl}(1)$:

where $p$ is the number of mantissa bits.

Dynamic range is determined by the exponent bits:

Key trade-off:

• BF16 sacrifices precision (only 7 mantissa bits vs. FP16’s 10) but preserves the full FP32 dynamic range (8 exponent bits). This is critical for training stability because gradient magnitudes can span many orders of magnitude.
• FP16 has better precision but a much narrower dynamic range ($\sim 10^{\pm 4.8}$), making overflow/underflow more likely.

The number of representable values between consecutive powers of 2 (e.g., in $[1, 2]$) is $2^p$:

6.3 FP16/BF16 Mixed Precision Training

Naively replacing all FP32 tensors with FP16/BF16 causes training divergence due to three failure modes, each addressed by a specific technique:

Trick 1: FP32 Master Copy of Weights

Problem — Weight update underflow: If a weight $w$has magnitude$\sim 1$and the gradient-based update$\Delta w$has magnitude$\sim 10^{-5}$, then in FP16:

because $\Delta w < \epsilon_{\text{FP16}} \times |w| \approx 10^{-3}$. The update is lost entirely. Once weights reach zero through underflow, they remain zero permanently (no gradient signal).

Solution: Maintain a FP32 master copy of weights $\mathbf{w}^{(32)}$. The training loop becomes:

Trick 2: Loss Scaling

Problem — Gradient underflow: Gradients are often much smaller than 1 (e.g., $|\mathbf{g}| \sim 10^{-6}$), falling below the minimum representable value in FP16 ($\sim 5.96 \times 10^{-8}$) or being too imprecise in BF16.

Solution: Scale the loss before the backward pass and unscale gradients afterward:

By linearity of differentiation, scaling the loss by $\alpha$scales all gradients by$\alpha$, shifting them into a representable range. The unscaling restores the correct values before any further processing (clipping, optimizer step).

Dynamic loss scaling starts with a large $\alpha$ and halves it whenever overflow (Inf/NaN) is detected in gradients, doubling it after a fixed number of successful steps.

Trick 3: FP32 Accumulation

Problem — Accumulation error: When summing many small values (e.g., computing means, batch norms, reductions), the running sum grows large while individual addends remain small, causing catastrophic cancellation:

The relative error of naive summation in precision $\epsilon$ is bounded by:

For $N = 4096$and$\epsilon_{\text{BF16}} \approx 7.8 \times 10^{-3}$: relative error $\leq 32$, which means the result can be completely wrong.

Solution: Accumulate intermediate sums in FP32 even when inputs/outputs are in 16-bit:

This is typically implemented at the hardware level in Tensor Cores, which accept FP16/BF16 inputs but accumulate in FP32 internally:

6.4 Per-Element Memory Cost Summary (Mixed Precision)

The full mixed precision training memory per parameter is:

6.5 FP8 Pretraining

FP8 matrix multiplications on H100 achieve twice the peak throughput of BF16:

However, FP8 introduces severe stability challenges due to the extremely limited dynamic range and precision.

Quantization for FP8

Converting a tensor $\mathbf{x}$ from high precision to FP8 requires scaling to fit within FP8’s representable range:

where $x_{\max}^{\text{FP8}} = 448$for E4M3 or$57344$ for E5M2.

DeepSeek-V3’s tile-wise quantization computes separate scaling factors per tile to reduce the impact of outlier values:

• Activations/inputs: tiles of size $1 \times 128$ (per-token granularity)
• Weights: tiles of size $128 \times 128$

For a tile $\mathcal{T}$:

This is much more robust than per-tensor scaling because a single outlier value in one tile doesn’t compress the dynamic range of all other tiles.

FP8 Mixed Precision Configurations

The primary stability risk of FP8 pretraining is that the E4M3 format has only $2^3 = 8$representable values per binade and a dynamic range of only$\sim 10^{\pm 2.4}$. Loss divergence typically manifests as:

1. Gradient underflow → zero gradients → stalled learning
2. Activation overflow → NaN propagation → loss explosion
3. Accumulation error → incorrect weight updates → slow divergence

The instability worsens with higher learning rates because the gradient magnitudes increase, exceeding FP8’s representable range more frequently:

6.6 Computational Throughput Benefit

The throughput gain from lower precision comes from two sources:

1. Higher peak FLOPS: Tensor Cores execute lower-precision operations in fewer cycles:

2. Reduced memory traffic: Smaller data types reduce HBM bandwidth requirements:

For memory-bound operations, this directly translates to a $2\times$ speedup. For compute-bound operations (GEMMs), the peak FLOPS doubling provides the speedup.

7. Connecting All Concepts: The GPU Performance Model

The overall throughput of a training step on a single GPU is determined by the roofline model:

where:

• $\Phi_{\text{peak}}$: peak compute throughput (FLOP/s) at the chosen precision
• $\text{AI} = \frac{\text{FLOPs}}{\text{Bytes transferred}}$: arithmetic intensity of the operation
• $\beta_{\text{HBM}}$: HBM bandwidth (bytes/s)

The ridge point — the AI where the transition from memory-bound to compute-bound occurs — is:

For H100 in BF16:

For H100 in FP8:

Operations with $\text{AI} < \text{AI}_{\text{ridge}}$(attention, LayerNorm, activations) benefit most from kernel fusion and FlashAttention (which reduce bytes transferred). Operations with$\text{AI} > \text{AI}_{\text{ridge}}$(large GEMMs) benefit most from lower precision (which increases$\Phi_{\text{peak}}$).

This unified view explains why all three techniques — kernel fusion, FlashAttention, and mixed precision — are complementary and simultaneously necessary for maximizing GPU utilization in modern large-scale training.