Welcome to the grand finale! In Part 1, we learned how to walk (GEMM). In Part 2, we learned how to run (backprop & AdamW). In Part 3, we learned how to fly across thousands of GPUs (distributed training). But if you are still here, you aren't satisfied with just "working" code. You want perfection. You want to squeeze every last drop of FLOPS out of your H100 or MI300X.
In this fourth and final part, we stop treating the GPU as a generic processor and start treating it as a memory-bound machine. We will implement:
- Flash Attention – The IO-aware kernel that made 100k+ context windows possible.
- FP8 Quantization (Transformer Engine) – Using 8-bit floating point for 2x speedups.
- Kernel Fusion – Fusing LayerNorm, residual adds, and dropout into a single pass.
- Distributed Checkpointing – Saving massive sharded models without crashing your filesystem.
By the end of this part, your custom training loop will rival the performance of PyTorch 2.0 + DeepSpeed. You will truly be a GPU wizard.
Prerequisites: The first three parts under your belt, a GPU with compute capability 8.9+ (Ada Lovelace/Hopper) or AMD CDNA 3 (MI300) to run FP8 natively, and a thirst for extreme performance.
- Flash Attention – Defeating the Memory Bottleneck
Standard attention computes S = Q * K^T (saving S to HBM), then reads S to compute softmax, then reads softmax to multiply by V. This means multiple round-trips to slow global memory (HBM) . Flash Attention solves this by fusing the entire attention pass into a single kernel using tiling and online softmax math.
The Core Trick: Online Softmax
Instead of storing the full attention matrix S (shape [seq_len, seq_len]), we process Q, K, V in small blocks that fit into Shared Memory (SRAM) . We recompute the softmax normalization on the fly.
Here is a simplified Flash Attention kernel skeleton (single-head, causal mask omitted for brevity):
#define TILE_SIZE 32 // Blocks of 32x32 in shared memory
__global__ void flash_attention_kernel(const float* Q, const float* K, const float* V,
float* O, int N, int d_head) {
// Shared memory tiles
__shared__ float Q_tile[TILE_SIZE][d_head];
__shared__ float K_tile[TILE_SIZE][d_head];
__shared__ float V_tile[TILE_SIZE][d_head];
int tx = threadIdx.x, ty = threadIdx.y;
int batch_idx = blockIdx.z; // assuming batch is in z-dim
int q_start = blockIdx.y * TILE_SIZE;
int kv_start = blockIdx.x * TILE_SIZE;
// Load Q tile into shared memory (per block)
for (int i = ty; i < d_head; i += blockDim.y) {
Q_tile[tx][i] = Q[((batch_idx * N) + q_start + tx) * d_head + i];
}
__syncthreads();
// Local accumulators for attention output
float out_acc[d_head] = {0.0f};
float l = 0.0f; // normalization sum
float m = -INFINITY; // running maximum
// Loop over all KV blocks
for (int kv_block = 0; kv_block < N / TILE_SIZE; ++kv_block) {
// Load K and V tiles into shared memory (load from global)
load_kv_tile(K, V, kv_block, ...);
__syncthreads();
// Compute Q * K^T for this tile (results in registers)
float scores[TILE_SIZE] = {0.0f};
for (int i = 0; i < TILE_SIZE; ++i) {
float sum = 0.0f;
for (int k = 0; k < d_head; ++k) {
sum += Q_tile[ty][k] * K_tile[i][k];
}
scores[i] = sum * rsqrtf((float)d_head);
}
// Online softmax update (for this tile)
for (int i = 0; i < TILE_SIZE; ++i) {
float score = scores[i];
float new_m = fmaxf(m, score);
float exp_diff = expf(m - new_m);
float exp_score = expf(score - new_m);
// Correct previous accumulated values
for (int k = 0; k < d_head; ++k) out_acc[k] *= exp_diff;
l = l * exp_diff + exp_score;
m = new_m;
// Accumulate weighted V
for (int k = 0; k < d_head; ++k) {
out_acc[k] += exp_score * V_tile[i][k];
}
}
__syncthreads();
}
// Final normalization and write to global O
for (int k = 0; k < d_head; ++k) {
int idx = ((batch_idx * N) + q_start + ty) * d_head + k;
O[idx] = out_acc[k] / l;
}
}
Note: This is significantly faster than Part 1's naive attention because we never write S to HBM. In practice, you use Flash Attention 2 or 3 (which adds warp-level parallelism), but this kernel teaches you the exact principle.
- FP8 Training (Transformer Engine)
FP8 (8-bit floating point) doubles the bandwidth and compute throughput compared to FP16. NVIDIA's Transformer Engine and AMD's FP8 support use a per-tensor scaling factor to avoid overflow.
We define a custom half-precision type (or use __hip_fp8 on AMD / __nv_fp8 on CUDA). Here is a cast kernel that scales a tensor down to FP8:
// Simplified FP8 cast with dynamic scale
__global__ void cast_fp32_to_fp8_kernel(const float* input, __nv_fp8_e4m3* output,
float scale, int n) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= n) return;
float val = input[idx] * scale; // scale up to use FP8 range
// Saturate to FP8 max range (roughly ~448 for E4M3)
val = fminf(fmaxf(val, -448.0f), 448.0f);
output[idx] = __nv_fp8_e4m3(val);
}
During training, we maintain a history of maximum values to dynamically adjust the scale per layer. For MatMul, we cast Q, K, V to FP8 before the main GEMM, and accumulate the result in FP32 (using hipblasGemmEx with compute type HIPBLAS_COMPUTE_32F).
- Fusing Residuals, LayerNorm, and Dropout
In a standard transformer block, we do:
- X = Attention(X) + X (Residual add) → Write to HBM.
- X = LayerNorm(X) → Read from HBM, Write to HBM.
- X = Dropout(FFN(X)) + X → Read/Write to HBM.
These are memory-bound operations. We can fuse them into a single kernel that reads once and writes once. Here is a fused kernel for the residual + LayerNorm path:
__global__ void fused_residual_layernorm_kernel(float* X, const float* Attn_Out,
float* gamma, float* beta,
int rows, int cols) {
extern __shared__ float sdata[]; // shared memory for reduction
int row = blockIdx.x;
int tid = threadIdx.x;
float* x_row = X + row * cols;
float* attn_row = Attn_Out + row * cols;
// 1. Compute residual add and mean/variance in registers
float sum = 0.0f, sq_sum = 0.0f;
for (int i = tid; i < cols; i += blockDim.x) {
float val = x_row[i] + attn_row[i]; // residual connection
x_row[i] = val; // store temporarily (we will normalize in-place)
sum += val;
sq_sum += val * val;
}
// Shared memory reduction (omitted for brevity, use warp shuffle)
float mean = sum / cols;
float variance = sq_sum / cols - mean * mean;
float inv_std = rsqrtf(variance + 1e-5f);
// 2. Apply LayerNorm in-place and write back
for (int i = tid; i < cols; i += blockDim.x) {
float normalized = (x_row[i] - mean) * inv_std;
x_row[i] = normalized * gamma[i] + beta[i];
}
}
This kernel does the work of 3 separate CUDA/HIP calls, saving two full read/write passes to global memory. That is roughly a 30% speedup for the non-GEMM parts of the model!
- Distributed Checkpointing – Saving Sharded Models
When training with ZeRO-3 in Part 3, each GPU only holds a fraction of the weights. If you naively gather all weights to rank 0 to save a single file, you will likely run out of memory and create a massive I/O bottleneck.
Instead, we implement Parallel Checkpointing. Each GPU saves its local shard (e.g., model_shard_rank_0.bin, model_shard_rank_1.bin) to the filesystem simultaneously.
void save_checkpoint_sharded(void* local_weights, size_t local_size, int rank, const char* base_dir) {
char filename[256];
snprintf(filename, sizeof(filename), "%s/checkpoint_step_%d_rank_%d.bin", base_dir, step, rank);
// Use buffered, asynchronous file writes to not stall the GPU
// 1. Copy from device to pinned host memory (async)
float* host_buffer;
hipHostMalloc(&host_buffer, local_size, hipHostMallocDefault);
hipMemcpyAsync(host_buffer, local_weights, local_size, hipMemcpyDeviceToHost, stream);
hipStreamSynchronize(stream);
// 2. Write to disk using standard fwrite or POSIX (on a separate thread ideally)
FILE* fp = fopen(filename, "wb");
fwrite(host_buffer, 1, local_size, fp);
fclose(fp);
// 3. Also save a metadata file (json) containing the world_size, shapes, and dtypes
hipHostFree(host_buffer);
}
To load, we simply reverse the process: each rank loads its own .bin file. This scales linearly with the number of GPUs and eliminates the "rank 0 bottleneck" completely.
- The Ultimate Asynchronous Data Loader (CPU Prefetch)
A slow CPU data loader will starve your massive GPU cluster. We use Double Buffering with pinned memory:
// Initialize two buffers
float* h_buffers[2];
float* d_buffers[2];
hipHostMalloc(&h_buffers[0], batch_size, hipHostMallocWriteCombined);
hipHostMalloc(&h_buffers[1], batch_size, hipHostMallocWriteCombined);
hipMalloc(&d_buffers[0], batch_size);
hipMalloc(&d_buffers[1], batch_size);
int current = 0;
for (int step = 0; step < total_steps; ++step) {
// While GPU processes buffer 'current', CPU loads buffer '1-current'
std::thread loader(load_next_batch, h_buffers[1-current]);
hipMemcpyAsync(d_buffers[current], h_buffers[current], batch_size,
hipMemcpyHostToDevice, compute_stream);
// ... Run Forward/Backward using d_buffers[current] ...
loader.join();
current = 1 - current;
}
This ensures your GPU compute kernel never waits for data. This is standard in production systems but often overlooked in custom C++ trainers.
Conclusion
And there you have it. We started from vec_add in Part 1, and now we have built a production-grade, multi-GPU, FP8-mixed-precision trainer with fused kernels and asynchronous everything. You have implemented Flash Attention logic, scaled gradients across nodes with Ring-AllReduce, sharded memory with ZeRO, and optimized memory traffic with Kernel Fusion.
Writing a custom trainer like this is a monumental task—which is why frameworks like PyTorch exist. But the next time you see a torch.compile warning, a DeepSpeed configuration file, or a Flash Attention import, you won't see black magic. You will see the exact C++/HIP logic we just built together.
You have officially graduated from "GPU user" to "GPU architect".
A final challenge for you: Try combining the FP8 casts with the Flash Attention kernel. It will break, it will be frustrating, and when you fix it, you will have a kernel that is faster than 99% of the open-source implementations out there.
If you want a Part 5 (though I thought this was the end!), we could explore Graph Compilation (static computation graphs) or CPU Offloading for when you run out of VRAM. Let me know in the comments!
Thank you from the bottom of my silicon heart for joining me on this journey. Keep your kernels tiled, your wavefronts full, and your HBM bandwidth saturated.
Until next time, happy hacking! 🚀
See ya on the other side of the LLM!
Top comments (0)