Multi-GPU Kernels, Intelligence per Watt, Heterogeneous Inference, and More | YC Paper Club

The speaker observes a growing trend of specialization at the chip level, driven by increasing demand and the need for tailored hardware solutions. He highlights that training data centers require different specifications than inference data centers; training doesn't necessitate high bandwidth, whereas inference does. This shift is exemplified by Google’s TPU8 Zebra Fish and Sunfish, marking a move towards more specialized ASICs.

Brennan points out a significant bottleneck in reinforcement learning (RL) simulations: the environment.step function is often still executed on CPUs, hindering rollout speed and making classic RL challenging. This highlights an opportunity to leverage GPUs for this computationally intensive task, although it's currently framed as 'weird GPU usage'.

The speaker introduces the concept of batch size one inference and the challenges it presents. Prioritizing latency in applications like voice agents (e.g., phone calls) necessitates single-instance processing, which can be computationally expensive and resource-intensive, leading to delays if not optimized for specialized hardware.

For training models, the primary metric is time per step, and location flexibility is key. Training data centers can be located far from users, even orbiting the sun, as long as efficient interconnects facilitate gradient communication. This contrasts with inference, which requires more localized processing for low latency.

The speaker emphasizes that selecting the appropriate transfer mechanism—whether it's TMA or the copy engine—is vital for optimal performance. For common multi-GPU operations like all-gather gem, TMA proves more efficient than relying solely on the copy engine. This choice significantly impacts data throughput and overall workload efficiency based on message size.

The discussion introduces two primary strategies for scheduling inter-GPU communication with intra-GPU work: intraSM overlapping and interSM overlapping. IntraSM involves dedicating threads within a SM to either computation or communication (warp specialization), while interSM assigns entire SMs to communication, sacrificing some compute resources. Each approach presents trade-offs regarding performance and resource utilization.

The Parallel Kittens (PK) framework builds upon previous work like Thunder Kittens, providing a set of CUDA programming primitives designed to simplify multi-GPU kernel development. PK abstracts away low-level complexities by offering data structures for each layer in the memory hierarchy (registers, shared memory, L2 cache, global memory) and communication primitives that automatically select efficient transfer mechanisms.

The speaker highlights the impressive performance of Parallel Kittens (PK), noting that with just 50 to 100 lines of device code, it can match or surpass hand-optimized kernels often containing hundreds or even thousands of lines. This demonstrates PK's ability to significantly reduce development effort while maintaining high levels of efficiency.

Stuart discussed the importance of quantifying the impact of increased 'intelligence per watt,' specifically mentioning potential effects on GDP and wages. He indicated that a forthcoming release would detail these metrics, suggesting a focus on translating computational advancements into tangible economic benefits. This highlights a shift towards measuring not just performance but also real-world utility.

The team is developing 'Open Jarvis,' a project aimed at operationalizing insights and enabling personal AI coding stacks to run directly on devices like laptops and workstations. This initiative seeks to reduce reliance on cloud services such as OpenAI's LMS, offering users greater autonomy and potentially lower costs by shifting processing power locally.

Mark outlined a spectrum of GPU kernel programming languages, ranging from mathematical expressions to BTX/inline assembly. He categorized these options based on trade-offs between performance and productivity, noting that developers must choose the tools best suited for their priorities. This includes established libraries like PyTorch, Triton, CUDA, as well as newer approaches like ThunderCons, Cutless, and QDSL.

Mark introduced 'KernelBot,' a platform designed to analyze GPU kernel submissions. Users can submit kernels written in various libraries, allowing for data-driven insights into the popularity and usage patterns of different programming tools. This aims to provide valuable feedback on which approaches are most effective within the community.

The speaker notes that kernelms exacerbate AMDall's law effects, meaning they tend to expose bottlenecks within a system. These bottlenecks can manifest as slow compilation times or any situation involving waiting, highlighting the need for optimization strategies. The discussion emphasizes the importance of addressing these inefficiencies to improve overall performance and reduce latency.

The speaker observes a trend where GP mode has effectively become 'pay-to-win' due to the extended test time scaling, often spanning days or weeks. This implies that those with greater resources (compute power and time) have an unfair advantage in testing and validation, creating a barrier for smaller teams or individuals.

Misha introduces Marlo, a startup focused on building workload-optimized heterogeneous infrastructure. The core premise is that inference isn't uniform; different phases (prefix cache lookup, prefill, KV cache creation, autoregressive decode) demand varying resources like compute, network bandwidth, storage, and memory. Therefore, co-designing systems with specialized hardware for each phase is crucial for efficiency.

Misha explains the concept of arithmetic intensity – the ratio of floating-point operations (flops) to data movement. He introduces the roofline model, originally developed for HPC, which helps determine whether a kernel is compute-bound or memory bandwidth-bound based on this ratio and machine characteristics. Understanding this relationship is key to optimizing performance by targeting bottlenecks.

The discussion highlights that memory bandwidth often acts as a significant bottleneck in GPU performance. To overcome this, the speaker proposes introducing 'System B,' designed to efficiently handle latency-bound kernels and memory bandwidth limitations. This involves offloading specific tasks, such as MA kernels, to SRM (Spatial Reconfigurable Machine) machines which offer low latency due to their on-die architecture, extending the interactivity lifespan of the GPU.

The speaker discusses speculative decoding as an application for heterogeneous systems. In this process, a 'drafter' generates potential token sequences which are then verified in parallel by a 'verifier.' By running the drafter on System B (optimized for speed) and the verifier on System A, it is possible to achieve faster forward progress and improved overall latency compared to running both components on System A alone.

The speaker emphasizes that designing heterogeneous infrastructure is a full-stack problem with implications beyond just hardware. This includes data center considerations like power density and cooling, as well as networking challenges to avoid bottlenecks. Performance modeling through simulation is crucial for guiding design choices and ensuring efficient operation of the combined systems.

The speaker introduces the concept of using game engines as learning environments, particularly for roboticists and self-driving car researchers. However, traditional game engine implementations are inefficient for high-throughput training workloads due to resource contention between parallel copies. The solution lies in 'batch simulators,' a single game engine simulating multiple instances concurrently.

Brandon suggests that there's a significant opportunity to create abstraction layers in GPU programming. These abstractions should provide good default performance while utilizing raw GPU horsepower, ultimately simplifying the process of GPU programming. He draws an analogy to Python, envisioning something similar for GPUs – a scripting language that abstracts away much of the complexity.