This article introduces a high-throughput "megakernel" for tensor-parallel inference with the Llama-70B model on H100 GPUs. Building on a previous low-latency design, the authors adapt their instruction-and-interpreter model to a different workload characteristic. The core optimization strategy is to aggressively overlap compute, memory, and cross-GPU communication to maximize the simultaneous use of different hardware resources. The article details several key techniques: intra-SM instruction pipelining to hide memory latency, a global work queue for dynamic scheduling across SMs to smooth out jitter, and interleaving different instruction types to overlap compute-bound and network-bound operations. A novel approach to the tensor-parallel forward pass, involving a "distributed transpose" after attention, is also presented. When integrated into the "Tokasaurus" inference engine, the megakernel demonstrates a 22% end-to-end throughput improvement over SGLang on the ShareGPT benchmark. The authors have released the highly experimental research code.
Maximizing GPU Utilization: A Throughput-Optimized Megakernel for Llama-70B
Core Concept: The Megakernel for High Throughput: Hazy Research introduces a single, fused "megakernel" for Llama-70B inference, designed to maximize throughput. Unlike traditional systems that use hundreds of separate kernels, this approach aims to use all parts of the GPU (tensor cores, memory bandwidth, NVLink) simultaneously by aggressively overlapping different types of operations.
A Novel Parallelism Scheme: The implementation uses a "sequence parallel" approach but modifies the standard transformer block. They replicate the O-projection matrix on each GPU and replace the post-attention reduce-scatter operation with a custom "distributed transpose." This reduces network traffic by a factor of 8, making it easier to hide communication costs.
Technique 1: Overlapping Within an SM: The megakernel's on-GPU interpreter uses pipelining to begin loading weights and activations for the next instruction while the current instruction's compute is still running. This keeps the tensor cores constantly fed, providing a 2-6% performance improvement.
Technique 2: Overlapping Across SMs: A "global work queue" allows any available Streaming Multiprocessor (SM) to atomically pull the next instruction from a single list. This dynamic scheduling smooths out runtime variations between SMs, preventing bottlenecks and improving performance by up to 14% at large batch sizes compared to static scheduling.
Technique 3: Overlapping Across GPUs: Communication is handled by dedicated "storer" threads within the megakernel, allowing compute threads to proceed with other work. By "interleaving" different instruction types (e.g., compute-heavy matmuls and network-heavy all-gathers) in the work queue, the system can run both simultaneously on different SMs, improving throughput by over 6% at large batch sizes.
Benchmark Results: Integrated into the Tokasaurus inference engine, the megakernel was benchmarked against SGLang using 65,536 prompts from the ShareGPT dataset. The megakernel achieved a total throughput of 23,468 tokens/s, outperforming SGLang's 19,170 tokens/s by over 22%.
Code Release: The code is available but is explicitly described as unsupported "research code" that is highly sensitive to the specific compiler and hardware environment.