# Lecture 112: Production Megakernels for Real-World Inference

## Executive summary

The lecture details 'megakernels,' a novel approach for optimizing production inference by moving beyond traditional kernel-based GPU programming models. Instead of launching multiple small kernels with associated overheads (e.g., global synchronization, SRAM clearing), megakernels fuse the entire forward pass into a single, persistent kernel running on an on-device interpreter. This significantly reduces launch overhead and enables fine-grained scheduling by allowing cores to dynamically pick up slack from slower units, maximizing GPU utilization.

## Key takeaways

- Mega Kernels vs. Traditional Kernels: Traditional GPU programming involves launching many small, self-contained kernels (grids) sequentially, incurring overheads like global synchronization and SRAM clearing between launches. Mega kernels persist the entire forward pass into one kernel, eliminating these boundaries.
- Compiler Search and Optimization: The system uses a compiler search mechanism to decide whether to mega-kernelize an operation or leave it as a traditional kernel graph (e.g., using Kublo ops). The choice is based on minimizing wall-clock runtime, especially when arithmetic intensity is low.
- Symbolic Representation of Work: To handle dynamic shapes (e.g., sequence length $S$ and context length $P$), the system represents tensor shapes, extents, strides, and barriers symbolically using expressions. This allows the compiler to generate a single work queue that doesn't need rebuilding when runtime dimensions change.
- Dependency Tracking (Barriers): Fine-grained dependency tracking is managed using barriers, which are treated as tensors. The goal is to maximize the number of barriers by partitioning tensor dimensions by dependency, allowing consumers to start processing before all producers have finished.

## Technical details

- GPU Programming Models: Traditional models program a single thread with localized views, requiring explicit global synchronization between kernels. Modern approaches aim for lower-level control by directly scheduling work onto every core (e.g., 5x5x1 launch grid).
- Kernel Launch Overhead Mitigation: The overhead of host-launched kernels is extremely high (hundreds of microseconds to milliseconds). Using CUDA graphs helps by allowing multiple kernels in a Directed Acyclic Graph (DAG), but even DAG nodes incur launch overhead. Mega kernels eliminate this by running the entire pass in one kernel.
- Memory Pipelining: Persistent kernels allow for memory pipelining (e.g., loading next inputs while storing current outputs) across multiple output tiles, which is impossible when hard kernel boundaries exist.
- Execution Model: The system uses an on-device interpreter (conceptually a for loop with a switch statement) to execute the sequence of instructions within the single mega kernel. This allows cores to atomically pop off instructions from a global instruction queue in DRAM.
- Multi-GPU Scheduling: For multi-GPU environments, the compiler can utilize both Single Program, Multiple Data (SPMD) and Message Passing Model (MPMD). For tensor parallelization (TP), standard SPMD is used; for multi-node setups, pipeline parallelism is preferred.

## Practical implications

- Mega kernels are most beneficial for small models with many relatively small amounts of work, where the overhead between traditional kernels is significant.
- The approach allows for mixing mega-kernelized sections (e.g., transformer body) with separate optimized ops (e.g., LM head) to minimize overall runtime.
- By keeping all dependency tracking and scheduling on-device, the host CPU is not involved in the hot loop, enabling continuous high-speed re-execution of the forward pass.

## Topics

GPU Architecture, Compiler Optimization, Machine Learning Inference, Kernel Fusion, Hardware Acceleration, Kublo, Hazy Research

Source: https://www.youtube.com/watch?v=loZ4xQ5RZuU