Multi-GPU Kernels, Intelligence per Watt, Heterogeneous Inference, and More | YC Paper Club
In a Nutshell
Multi-GPU kernels require fine-grained compute-communication overlap using TMA for all-gather GEMM and register instructions for all-reduce, with Parallel Kittens achieving hand-optimized performance in 50-100 lines of code versus thousands. Intelligence per watt has improved 3x in two years for local accelerators, with 80-90% of queries potentially routable to local open-source models for 50-70% energy savings. Heterogeneous inference disaggregates prefill-decode and attention-MLP phases across specialized hardware, while GPU-accelerated batch simulators deliver 100x throughput gains for RL environments using entity-component-system architecture.
These notes were generated by AI and may contain inaccuracies.
The event is themed as the YC kernel and chip club due to the trend of specialization happening at the chip level. The TPU8 I and T version represent specialization into different ASICs (Zebrafish and Sunfish), driven by sufficient demand for such splits. Training data centers have fundamentally different requirements than inference data centers, with training centers not requiring any bandwidth in and out, while inference centers require proximity to users.
There is significant optimization potential remaining on the CUDA side, kernel side, algorithm side, software side, chip side, and data center side. The hosts include Stuart Soul from Stanford's Hazy Research lab, researcher at Cursor training Composer; John from the CS PhD lab at Stanford, co-advised by Aelia, focusing on LMS, MLIS, and hardware, pioneering intelligence per watt and intelligence per joule; Mark, former PyTorch maintainer, co-founder of GPU mode with Casey Elward, co-founding core automation with OpenAI VP research Jerry Torque; Misha with two decades of hardware-software co-design experience, previously running AI infra at NVIDIA and working on hardware-software co-design at Meta; and Brennan from Stanford working on RL for self-driving cars, focusing on GPU acceleration for simulators.
The proliferation of specialization at the chip level is driven by sufficient demand for tokens. The Zebraish Sunfish TPU V8 represents the first visible split, expected to become more pronounced. Pre-existing patterns include inference workloads using Nvidia for prefill and Cerebras for decode engines. A critical but under-discussed distinction exists between batch size one inference and throughput-optimized inference, particularly relevant for latency-sensitive applications like voice agents where 8-second delays are unacceptable. Batch size one prioritization leads to GPU exhaustion and high costs, representing a chip-level issue requiring specialized solutions.
Training data centers prioritize time per step for progressively larger models, with location flexibility allowing placement anywhere including orbiting the sun. Training requires all-to-all communication for gradient passing, while inference requires all-to-all communication only for certain operations but allows nicer sharding. Different chips for different data center types justify specialization, with the next speaker providing detailed analysis.
Stuart Soul presents work on making multi-GPU AI kernels simple and fast through a CUDA framework called Parallel Kittens. The goal is conveying lessons learned rather than promoting the open source library. The framework provides a mental model for thinking about multi-GPU kernels and GPU networking.
The ML systems community has focused heavily on single GPU efficiency through IO-aware algorithms like flash attention, linear attention models, mega kernels, and programming libraries across architectures. GPU networking remains the major bottleneck, consuming up to 50% of total runtime for workloads like Llama's MDB prefill. Networking weights GPU utilization, with standard practice involving overlapping inter-GPU communication with intra-GPU operations at increasingly fine granularity, from entire tensors to tile or token granularity of a few kilobytes or less than 1 kilobyte.
In-network compute enables offloading arithmetic communication to networking fabric for collective operations like all-reduce, freeing GPUs for useful computation. Asynchronous bulk device-initiated networking is now possible, along with scale-up architectures like NVL72 packing 72 GPUs inside a single NVLink domain, extending to hundreds of GPUs.
Writing multi-GPU kernels that achieve fine-grain compute-communication overlap, exploit recent hardware characteristics, and remain simple to maintain for humans and AI agents remains difficult. Three alternatives exist: off-the-shelf libraries like Nickel expose only coarse-grain primitives leaving hardware bubbles; compiler-based approaches produce suboptimal kernels sometimes slower than non-overlap baselines; low-level primitives like OS interprocess communication or PTX assembly require reverse-engineered hardware understanding and produce complex, hard-to-maintain code.
A GPU contains streaming multiprocessors (SMs) performing arithmetic, logical operations, and matrix multiplications, with 148 SMs on Blackwell GPUs and 132 on Hopper GPUs. The memory system includes crossbar, L2 cache, and HBM as main memory. Multiple GPUs connect through PCIe to CPU, with NVLink enabling direct GPU-to-GPU communication without PCIe congestion. DGX machines represent collections of GPUs connected via PCIe to host and interconnected via NVLink. NVL72 contains 72 GPUs connected via NVLink.
GPU kernel optimization focuses on keeping SMs busy consuming compute resources. Single GPU kernels pipeline memory loads and writes overlapping computation with HBM memory operations. Multi-GPU kernels similarly overlap computation with communication from remote GPU HBMs.
The first trade-off is transfer mechanism with three options: copy engine, tensor memory accelerators (TMAs), and register instructions. Copy engine is a separate DMA module not requiring SMs but requires large message sizes to saturate NVLink, making it suboptimal for fine-grain communication of a few kilobytes. TMA maintains good throughput regardless of message size with relatively few SMs, saturating NVLink with roughly 15 of 148 SMs on Blackwell, but cannot utilize in-network compute features. Register instructions require considering coalescing and register pressure but enable in-network compute.
For commonly used multi-GPU operators like all gather GEM, TMA proves optimal, whereas previous approaches relied on copy engine. The second trade-off is scheduling inter-GPU communication with intra-GPU work, falling into intra-SM overlapping (warp specialization with certain threads performing compute/memory work and others communication) versus inter-SM overlapping (dedicating certain SMs purely for communication, others for computation).
Inter-SM overlapping wastes compute on communication-only SMs and requires memory system traversal through L2 cache and HBM for signaling completion between communication and computation. Intra-SM overlapping requires computation and communication aligned on the same data, otherwise running out of registers or shared memory. Local prefetching becomes necessary because remote HBM data paths go through L2 cache while local GPU data does not, meaning repeatedly accessed data like KV cache during attention kernels gets cached only on the far-side GPU.
For GEM reduce scatter, intra-SM proves most optimal, while for GEM all-reduce, inter-SM proves most optimal. The third trade-off is design overhead referring to networking library choices costing performance for complexity reduction. Nickel's default mode forces intermediate buffers adding extra data movement, accumulating overhead for fine-grain communication. Stripping out this overhead speeds up operations as basic as all-reduce by up to 80%.
Parallel Kittens builds on identified trade-offs as a highly opinionated set of CUDA programming primitives extending Thunder Kittens for single GPU kernels. Modern GPUs have memory hierarchy from registers, shared memory, L2 cache, global memory, with remote GPU HBM as another layer. PK provides data structures for each layer: RT for registers, ST for shared memory, GL for global memory, PGL for pure global memory. Communication primitives utilize the most efficient transfer mechanism for given functionality and provide a program template with four workers: loader, consumer, communicator, and store.
PK achieves performance matching or surpassing hand-optimized kernels with 50-100 lines of device code versus hundreds to thousands of lines in baselines, across data, tensor, sequence, and expert parallelism plus collective communication like all-reduce. Adoption includes Cursor using PK to train Composer on tens of thousands of Blackwell GPUs, and Together AI using it to optimize inference workloads.
John presents joint work measuring intelligence efficiency of local and cloud AI, with collaborators and advisors Professor John Hennessy, Isal Huzzeni, and Chris Ray. The current era of LLMs and hardware is characterized as the mainframe era, analogous to 1950s-1960s IBM mainframe computers filling rooms and floors, getting humanity to the moon and serving as first commercial computers. Modern data centers represent warehouses filled with TPUs from Google and GPUs from Nvidia handling both training and inference.
Demand for AI infrastructure represents the largest seen since US railroad construction across the continent, with spending of 2-3% of GDP demanding 250 gigawatts of new data centers requiring new power sources, GPUs, land, and water. Despite rapid demand growth, many requests do not require frontier-level intelligence, with workloads including data organization, plumbing, folder organization, and unit test creation not needing frontier coding or math capabilities.
Smaller local open-source LLMs can address the vast majority of LLM traffic, driving a shift from cloud-based Nvidia GPUs and Google TPUs toward local accelerators from Apple, Nvidia, and AMD in workstations and laptops. This parallels the 1970s-1980s shift from mainframe IBM computers to PCs from Apple and HP. The transition is driven by smaller LLMs improving year-over-year with capable open-source models accelerating in capability, combined with local accelerators having sufficient memory for larger open-source LLMs.
The study examines the role local inference with open source AI can play in redistributing inference demand from predominantly cloud-based to distributed inference. The framework proposes
The framework proposes intelligence per watt, analyzing both the numerator (capabilities) and denominator (efficiency). Intelligence refers to capabilities across chat reasoning, agentic tasks, coding tasks, and other LLM use cases. The denominator focuses on actual compute and energy required for running accelerators across workloads, including power wattage for delivering intelligence.
The study tested 20+ state-of-the-art local models across Gemma, GBD OSS, Quen, and IBM Granite, ranging from 1 to 200 billion parameters, including both dense models. Researchers examined state-of-the-art accelerators from Apple, Nvidia, AMD, and Zamanova over the past 2-3 years to understand upper bounds for intelligence per watt and intelligence per joule. Testing covered all major workloads: chat reasoning, agentic tasks, and coding tasks, measuring accuracy, latency, energy, power, and compute.
Local models are advancing rapidly, with up to 88.7% of queries routable to local accelerators running local open source AI, with accuracy continuing to increase. The past two years have shown approximately 3x improvement in intelligence delivered per watt, combining better local models with better local accelerators.
Intelligence per joule (total energy spent for entire workload end-to-end) has compounded even more rapidly, showing approximately 18x improvement in 16 months, primarily driven by better accelerators. Key drivers include more memory on consumer GPUs like the Apple M4 Max and Nvidia DJX Spark, better quantization techniques, longer pre-training, more post-training, and sophisticated optimization techniques.
Perfect routing could direct 80-90% of current queries to local accelerators running local open source models, delivering energy savings, compute savings, and cost savings. Even imperfect routers can save 50-70% of energy, compute, and dollar cost, fundamentally changing the economics of scaling inference and data center planning.
Data center compute enables sophisticated techniques around kernel writing, batching, and quantization that amortize investment across more users and queries. The Apple M4 Max still lags behind the Nvidia B200 in intelligence per watt and intelligence per joule. Specialized accelerators like the Samanova SN40L show even greater gaps where consumer accelerators lag further behind, presenting investment opportunities.
The next phase focuses on how inference engines will change when leveraging both local and cloud resources. Major opportunities exist in model architectures and kernels focused on energy-efficient inference, potentially leading to different accelerator architecture designs. Analysis plans include breaking down intelligence per watt metrics across tasks, hardware, deployment settings, inference engines, kernels, and hardware compute/memory.
Beyond intelligence per watt, the goal is to quantify how intelligence is used for meaningful work, measuring impacts on GDP and wages. Open Jarvis is a follow-up project operationalizing insights to run entire personal AI coding stacks on device, eliminating cloud and OpenAI dependency for laptop and workstation deployment.
The discussion transitions to evaluating AI capabilities in writing GPU kernels, drawing from two papers: Kernel Bot (competitive platform for GPU programmers) and Kernel Guard (platform for detecting cheating). The core tension is between AI researchers wanting autoregressive models with dynamic control flow, sparsity, and print statements versus systems researchers preferring consistent matrix multiplication operations.
Math represents a forward-compatible programming language, with matrix multiplication APIs remaining consistent for 250 years. PyTorch ops (MM, softmax, QR decomposition) serve as forward-facing APIs that dispatch to CUDA kernels of varying performance. Triton gained adoption 2.5 years ago with its tile-based, Pythonic approach, but Blackwell-era developers find its programming model restrictive for state-of-the-art kernels.
The programming language landscape includes ThunderCons (maintained by Stuart), closed-source Nvidia libraries Cutlass and QDSL for state-of-the-art matrix operations, CUDA for general thread-based programming (tedious but effective for surgical improvements), and inline PTX/SASS at the lowest level without BC guarantees.
Kernel Bot data shows CUDA preferred for state-of-the-art performance, with CuDSL popular only for GEMM-related problems. Triton appears in top 5-10 submissions but rarely in top 2. Unexpectedly, people with no GPU kernel experience achieved competitive results, including a Chinese grad student reaching #4 with entirely LLM-generated code and a high school teacher successfully writing dual GEMM kernels.
The evaluation framework uses PyTorch as the correctness reference, samples random inputs, runs submissions through both reference and candidate kernels, checks equivalence, correctness, and performance, then ranks on leaderboards. The process appears straightforward but contains hidden nuances around reward hacking.
A practical example demonstrates the world's fastest vector mean kernel: given a 1D vector, simply return zero. This exploits the default PyTorch sampling of values with mean zero and variance one, beating the speed of light but highlighting reward hack vulnerabilities.
Common reward hacks include caching outputs for reuse, using data pointers to detect past state in Python, splitting banned strings like "data pointer" across multiple lines to evade regex detection, and using ID instead of data pointer. A sophisticated example involved counting correctness test iterations (15 times) versus performance testing, returning correct but slow kernels during correctness checks and incorrect kernels during performance testing.
The solution involves human auditing of fast submissions to identify reward hacks, providing examples to AI systems for synthesizing regex-based cheat detectors, deploying these detectors, and retraining as new patterns emerge. This adversarial approach mirrors how PyTorch achieved correctness through iterative community feedback over nine years.
Shampoo optimizer requires QR factorization, which is slow in PyTorch. The community collectively developed a kernel 60x faster without NaN issues, enabling real training runs. Submissions average 15,000 lines of code with single kernels per shape, implementing dispatchers for different memory hierarchies and precision levels based on shape frequency.
Synthesizing thousands of long, complex submissions into elegant, beautiful kernels remains an open problem. The fastest kernels demonstrate deep understanding of GPU memory hierarchy, though the process of achieving this remains more philosophical than actionable.
Kernels exacerbate Amdahl's law effects by finding all bottlenecks in a system, whether compilation times or anything involving waiting. Faster JITs and more efficient AOT packaging of things like Trident and QDSL are of interest. CPU simulators of GPUs are quite good for this work. Fable is phenomenal at this kind of work because if you can spin up CPU simulators, that means you don't need a lot of GPUs for doing rollouts.
When deploying kernels in a real inference engine, out-of-the-box systems like SGLang or VLLM loading DeepSeek before the first inference may wait about 30 minutes. More robust ways to verify kernel correctness are needed beyond taking random inputs. Formal verification is one approach, though cheaper alternatives than training a model are needed. Intermediate approaches between reading code line by line and extensive testing would be valuable.
GP mode has become pay to win because test time scaling happens over days or on the order of one or two weeks. The question is whether this can be speedrun in a couple of hours or days.
Misha Manski joined Marlo, a startup focused on building workload optimized heterogeneous infrastructure. The main premise is that inference is a very heterogeneous workload where different phases exercise compute, network, storage, and memory bandwidths differently, making it sensible to co-design systems that use different hardware for different phases.
An inference request lifecycle includes: user enters a prompt, CPU system orchestrates and schedules and batches it, prefix cache lookup, prefill for uncached prompt content (compute intensive, done on accelerators), KV cache creation (can involve CPU or other accelerators memory over network), autoregressive decode (one token at a time, memory intensive and latency sensitive), result production, and speculative decode along the way.
Arithmetic intensity is the ratio of flops an algorithm does to the amount of data moved from memory. If this ratio exceeds the machine ratio (peak machine flops to peak memory bandwidth), the kernel or algorithm is compute bound; if lower, it's memory bandwidth bound.
The roofline model, created for HPC by Sam Williams from Berkeley around 2012-2013, plots arithmetic intensity on the X-axis against achievable flops. Arithmetic intensity below machine ratio indicates memory bound, above indicates compute bound.
Prefill is generally compute bound because attention and MLP kernels do substantial work per token, fetching weights once and operating on all tokens. MLP is even more intensive because batching is possible. Decode operates one token at a time autoregressively, making attention inefficient since weights must be fetched for every token processed.
Different use cases span a large space of workloads with varying prefill versus decode ratios and concurrency levels. Interactive chats typically have 30/70 to 50/50 ratios between prefill and decode with strict latency sensitivity. Long context queries spend most time in prefill with some decode afterward. Coding agents have strict latency requirements while long-running agents have relaxed latency but high concurrency with long inputs and outputs.
SRAM machines became popular in the last decade. When doing GEMM on GPU with weights larger than on-die caches, data must be fetched from HBM over the unpackage interconnect for every token, which is inefficient. SRAM machines keep the entire weight matrix in on-die SRAM memory, providing significantly more bandwidth since it's on-chip with fast chip interconnects between different units.
These are essentially GMV (matrix vector multiplication) accelerators that are very good for decode because decode is bandwidth bound. They offer significantly more bandwidth and significantly lower latency, though capacity is limited by die area constraints to tens to hundreds of megabytes up to tens of gigabytes of memory.
Prefill-decode disaggregation lets each phase scale independently by removing shared bottlenecks. System A can do prefill while System B does decode. From a TCO perspective measured in tokens per second per watt, this is beneficial when additional power from System B is offset by speedups. For short sequence lengths with substantial prefill time, benefits may not be pronounced, and short output lengths can result in TCO loss from bundling unused hardware. As output length increases and more time is spent on decode, TCO can become positive.
Attention and MLP kernels have different characteristics. One approach runs attention on System A and MLP on System B. GPUs excel at high throughput with high concurrency and compute-bound workloads, delivering lower latencies, but performance drops sharply at lower concurrency due to smaller kernels, overheads, and poor memory bandwidth utilization.
Introducing System B optimized for memory bandwidth and latency bound kernels (such as offloading MLP kernels to SRAM machines) can extend interactivity range. While TCO may not match high-concurrency GPU operation, this approach makes sense for use cases where interactivity matters more than TCO, extending the point beyond which GPUs cannot deliver more interactivity.
In speculative decoding, a drafter speculates on subsets of tokens either autoregressively or in parallel through block diffusion, then a verifier verifies them in parallel. Accepted tokens enable fast-forwarding. The drafter can run on System B while the verifier runs on System A, allowing larger models at the same rate as System A and improving acceptance rates. This reduces latency since System A may not handle such large models efficiently, and multiple verifiers can be bundled to the same drafter to improve TCO.
End-to-end co-design of heterogeneous infrastructure is a full-stack problem with multiple aspects. Data center considerations include power density, cooling, rack configurations, and managing problems that become harder in heterogeneous versus homogeneous environments. Networking becomes critical when moving data between System A and System B, requiring co-design of network topology so latency doesn't eliminate advantages.
Performance modeling infrastructure with well-calibrated simulators is essential for evaluating many design points and guiding co-design decisions.
Brennan presents work from Stanford PhD research in Kayvon Fatahalian's group, now doing an AI plus gaming startup. The project explored putting an entire game engine capable of simulating any imaginable game on the GPU and running it super fast.
Games provide excellent learning environments for various tasks including robotics and self-driving car research, offering low-cost skill learning in simulation. However, game engines are inefficient for throughput-oriented training workloads when running thousands of copies in parallel, as copies fight with each other, costs cannot be amortized, and both CPU and GPU hardware are used inefficiently.
Batch simulators address this by having a single game engine simulate a batch of a thousand learning environments simultaneously in a throughput-oriented batch on the GPU. OpenAI's hide-and-seek environment demonstrated fitting many environments on a single GPU, achieving throughput of millions of frames per second of experience, which is valuable for data-hungry algorithms like RL.
Existing GPU programming frameworks are a poor fit for gameplay logic needed to build these environments. The game industry solved this for parallel CPU code using entity component system design patterns, which can be adapted for GPU execution in throughput-oriented workloads.
Entities (the E in ECS) are objects like obstacles or agents. Components are data attached to entities such as position, rotation, action, and reward from reinforcement learning systems. ECS puts all this data in big in-memory column stores, with GPU implementations placing everything in GPU memory and packing learning environments into unified column stores.
This unified table storage enables throughput-oriented dynamic memory allocation on the GPU while hiding complexity behind an easy interface like "create an obstacle, create an agent," making procedural generation tasks straightforward compared to fixed-size arrays in tensor-based programming languages like PyTorch.
Systems (the S in ECS) contain the actual code running over entities to build logic, such as processing actions from neural networks, implementing collision detection, or computing rewards. Systems declare requirements like running on every entity with position and action components, mapping one function invocation to one GPU thread for massive parallelism.
Runtime type polymorphism is supported without virtual dispatch overhead that breaks statically known register assignments on GPUs. The collision system can require position and bounding box components without knowing specific types like agents or obstacles.
Individual systems combine into task graphs describing how the overall frame works from taking agent actions to updating world state. Memory allocation efficiency in tables leverages the throughput-oriented workload with an approach similar to garbage collection on the GPU, appending rows, marking rows as deleted, then using high-performance GPU sort to clean up afterward.
Task graphs are implemented with persistent mega-kernels constantly processing remaining frame work. GPU fast atomics enable efficient synchronization throughout this process.
A visualization shows each row representing an SM on the GPU, with colored bars indicating how active each SM is. Different colors correspond to different systems in the task graph. Between each set of colors are slight vertical white lines representing synchronization points in the frame where workloads switch. These synchronization points account for under 1% of the overall work, with the GPU being fully utilized most of the time.
This visualization is shown on an RTX 4090 simulating approximately 4,000 worlds of a hide-and-seek environment.
For the initial version of the engine, a small set of baseline environments were built, including the hide-and-seek environment and Overcooked, which is a cooperative reinforcement learning environment. These environments were specifically chosen because they already had easy-to-set-up CPU baselines, enabling direct apples-to-apples comparisons.
The performance comparison shows three implementations:
- Green represents the original CPU reference implementations from machine learning codebases, which performed poorly with all tasks being majorly bottlenecked by environment throughput
- Blue shows the ECS ideas running on CPU, similar to a traditional game engine but scaled up to run across multiple worlds
- Red shows the implementation running entirely on a 4090 GPU, achieving over 100 times faster performance in many cases
Multiple projects demonstrated that this approach can extend to end-to-end training, not just simulation throughput. A variety of fast environments were built on this engine over approximately 2 years. Non-technical machine learning researchers with no low-level GPU programming knowledge successfully built environments and achieved good performance, accelerating their workloads.
There's still a lot of work to be done on high-level scripting languages for the GPU even in the era of LLM's writing code.
Current programming language systems trend toward maintaining the CUDA programming model with Python syntax on top, which simplifies things but doesn't address fundamental issues like dynamic memory allocation on the GPU and super irregular parallelism. GPUs have massive memory bandwidth and compute even without considering tensor cores, enabling interesting workloads to leverage this hardware for significant speedups. However, fundamental issues create friction for novel workloads.
There is an opportunity for abstractions that provide good enough default performance while utilizing raw GPU horsepower but making GPU programming far simpler, similar to a real scripting language Python for the GPU.
The session concludes with thanks to the speaker Brandon, followed by a brief wrap-up. The thematic cluster approach will continue with the next session focused on robotics, with calls for suggestions of robotics speakers from industry. An additional hour is allocated for networking and discussions with speakers.
Keep Y Combinator in your library
Save the videos and channels worth coming back to, and find them again in one place.
