Can LLMs Write Fast Multi-GPU Kernels? — Simran Arora, Together AI

[music]

Hi everyone. Uh, sorry it's a bit loud

in here. Was not expecting this. Um, I'm

Siman. I'm a principal scientist at

Together AI. Um, I previously did my PhD

in the Hazy Research Lab with Chris Ray

at Stanford and I'm an incoming

professor at Caltech. Um I lead the

frontier performance research team at

together where we develop systems,

frameworks and algorithms to extract as

much performance as possible out of

modern um AI hardware. Today I want to

share a little bit about our

contributions towards simplifying the

development of uh multi-GPU AI kernels.

A few years ago, um, GPU utilization

used to be limited by poor intraGPU

memory access and single GPU kernels.

But with significant investment in

better kernels like flash attention, uh,

memory efficient architectures like from

deepseek, sparse attentions, mambas and

so on, um, and better DSLs, we've sort

of shifted the bottleneck to multi-GPU

communication.

During this talk, I'll start by telling

you a little bit about why now, why GPU

networking now. Then I'll tell you about

um the sort of problem space. So what

are the challenges in maximizing

hardware utilization and development

simplicity for multiGPU kernels. Um

three, we'll talk a little bit about the

fundamentals be behind designing

effective multiGGPU kernels. Four, we'll

look at whether frontier AI models can

uh leverage these fundamental

principles. Do they understand them? Can

they reason about them? Um, you know, in

theory, these models are very good at

reasoning. Um, and then five, we'll talk

through the results of these frontier

models on a benchmark that we've

developed called um parallel kernel

bench for multiGPU kernel generation

evaluation.

Okay, before we dive into those five

parts, just basic preliminaries. So,

this is an Nvidia GPU, uh, an H100 GPU

that you can see on the screen. Um, I

always like to help ground people in GPU

kernels via looking at the hardware. So,

these rainbow colored dots are

processors where actual compute is

happening. um all of the you know

parallel threads are operating within

one of those colored dots and there's

typically you know 100 200 of them on

modern AI GPUs.

Um around those processors you can see

um some of the memory that these

processors retrieve data so large

weights activations from. So these

rectangles between the colored dots are

an L2 cache slightly faster memory. um

not a ter like crazy large amount of it.

And then these black boxes are high

bandwidth memory. So when you Nvidia Smi

and see you know 80 gigabytes, 1008

gigabytes, whatever it is on your GPU,

that's that memory.

Um a GPU is operating a highly parallel

program. So multiple threads are

combined together in uh into larger

coarser units. and we schedule these

threads and and blocks onto these

processors to perform our AI compute.

Beyond the GPU, we'll have multiple GPUs

and we'll also have, you know, CPUs um

that have memory as well.

Um so to perform computation um the

memory that these threads use is going

to be stored in a really fast register

memory that's right next to the

computation units. um simple physics if

I am pulling data from very very close

to my proc my compute unit it's really

fast to get to it because it you know

that data is right next to me but

there's not a large radius and not a

large volume of space that's close by to

my process uh my my compute units and so

I don't have very much of it so you can

see that the fastest memory here the

registers is sup is you know 130

terabytes per second on an H100 but we

don't have very much on of it and as we

go to the further away memory we have a

lot more of it but it takes longer to

reach it. Uh again simple physics.

So in multiGPU systems in particular um

there is a hierarchy of interconnects.

So we will have something called PCIe um

as the channel for CPUGGPU

communications. We'll have multiGPU or

multi- uh node communications over

infiniband TCP and then in the we're

going to focus mostly on the Nvidia

sphere here. Um in uh the intraGPU

regime we'll have NVLink um providing

point-to-point connections between GPUs

and the NV switch. NV switch connects

all NVLink in endpoints into a

non-blocking fabric for full GPUGGPU

communication and NV switch is exciting

because it also provides support for in

network offdevice acceleration for oper

like communication primitives like

multiccast and reductions.

Okay, so diving in with the

preliminaries in mind. Why GPU

networking now?

So as I mentioned at the beginning,

we've really put a lot of effort into

making AIO uh more efficient over recent

years. Um again, we have uh

architectures that use less compute,

less memory like Mamba or sparse

attentions. We have algorithms that make

AI more hardware affair aware like flash

attention. We have tools to make it easy

to map AI algorithms to the hardware

like tileang, mojo, triton, gluon,

thunder kittens. Um, and we have new

techniques to overlap uh execution

across many AI operators very tightly

like mega kernels. And we also have

tools to make it easy to run on multiple

vendor and silicon platforms like

thunder mittens for Apple silicon or or

hipkittens for AMD and so on.

Um, at this point we really believe that

GPU networking offers many new and

exciting opportunities for AI

efficiency.

Modern AI workloads are getting very big

and require kernels that span multiple

um, you know, GPUs. So on many

production um distributed training and

inference workloads, communication is

increasingly consuming the majority of

the runtime and yields low model flop

utilization at scale.

So the uh pace of innovation and

diversity of approaches that different

hardware providers are taking in their

networking stacks is another reason why

it's an exciting time to study um you

know networking and communication. So we

can see here um AMD hardware with uh

what's called XGMI interconnects um

providing pointto-point links between

different GPUs in a scaleup domain. We

can see here TPU and uh interconnect uh

as well. So the TPU will use a 3D Taurus

and also have optical wraparound links.

So yet another diverse form of the

topology and links. Um and then again

for Nvidia we'll have the INV um switch

which integrates um compute capabilities

directly into the interconnect fabric

like for in network reductions and then

we'll have our NV link providing up to

you know 900 gigabytes of unidirectional

bandwidth between any two remote GPUs

high bandwidth memory um on you know

particular generation of Nvidia hardware

um Beyond the diversity in the

networking stacks, there's also a lot of

evolution in um how AI workloads are

adapting to take advantage of this

hardware and the diversity of types of

um you know hardware that we're we're

using simultaneously for one AI

workload. So um KV cache memory in

modern inference systems is going to

beworked across GPU, CPU, disk and

remote machines. um inference systems

increasingly disagregate different steps

of inference across different hardware

backends. So you could run speculative

decoding on some hardware, decode on

different hardware, prefill on different

hardware. Um hardware is also evolving

to have larger scaleup domains than ever

before with you know 72 GPUs and a

scaleup domain in in the coming um chips

and Nvidia planning on a single system

in 2027 with 576 GPUs.

Um at the same time to take advantage of

these more uh intensive scaleup domains

we're getting richer uh primitives for

fine grained control and kernel writing

over these domains. So we have something

called tensor memory acceleration where

we can provide uh perform asynchronous

network transfers from the device side

um on these GPUs.

So all of these changes are opening up

new opportunities and challenges in um

both AI you know how do we build models

that take advantage of these trends and

in systems.

Okay, so the problems that we're going

to go after um how do we get peak

hardware utilization and also

development simplicity um for these

multiGPU kernels.

So um it's been very difficult to write

multiGGPU kernels and there's a lot of

there are a lot of papers a lot of

systems reports that document you know

challenges here. Um, one of the things

here is it's compounded by the fact that

communication hardware around GPUs has

progressed a lot more slowly relative to

compute and memory. Um, so comparing

NVIDIA A100's in 2020 to B200s in 2024,

u BF16 tensor core speeds improved by

7.2x.

um while intra node communication by

just 3x and inter node communication by

just 2x. Um and coming back to my points

about how diverse networking is right

now things like tensor cores that run

map moles and our memory hierarchies are

pretty consistent and resemble one

another across diverse AI vendors and

multisilicon. Um but again as I I

mentioned the networking stack is

something that is really different

across vendors still

um you know a first step as we went

about all this work is to just study the

baselines um that are out there. So um

one of the popular tools for um

communications is this nickel library or

Rickle on AMD um that both you know

companies respectively spend a lot of

you know engineering investment into

releasing to make it easy for people to

do multiGPU work. Um but they're not

very flexible. So they're tuned for bulk

transfers for large contiguous chunks of

data transfers. And the design really

breaks down when you care about peak

performance, fine grain communication,

um, and sort of non-trivial collectives

that you want to fuse together.

So as a result, you can achieve much

higher performance by writing custom

communication kernels that directly

address these needs.

If we look at a naive baseline that's

representative of very popular libraries

in machine learning stacking pietorch

with nickel um we can we find that

across um you know the many uh problems

in our parallel kernel bench benchmark

that the majority of these um simple

baselines will fall below 50% of their

communication aware roof line bound. So

there's a lot of room for improvement

here.

Um the current frameworks beyond um

nickel which is popular in like Megatron

LM, Flex Flow, Nanoflow um you know all

again these systems are primarily

orchestrating bulk collectives via

nickel um and require synchronization

before and after data transfers. So

beyond th those off-the-shelf libraries,

we have compilers and DSLs that exist.

So there's Triton distributed is one of

them. Um and uh you know tile link is

another one. Um we have found that it's

very difficult to support the rapid pace

of networking improvements within these

frameworks. So our benchmarks and our

papers highlight results where Triton

distributed originally tuned around 8 uh

H800 GPUs fails to adapt efficiently to

other architectures like H100s.

And then the third category of how

people can proceed here is to really

handtune specific AI operators one by

one. So there's a lot of popular work um

DPP um comet ring attention um flux

flashdoe

um and then several distributed gem

kernels from cutless and these methods

achieve peak performance but often um

they do not you know some of these

methods have been designed in one

precision and it takes five or six

months to scale it to another precision

and just the scalability of this hand

tuning and um fine grain kernel writing

is not very um effective.

So um with this landscape in mind, our

research question was really about

whether there is a small set of

principles and fundamentals that really

governs multiGPU kernel writing and

whether there are methods that can

leverage those principles if they exist

to simplify the development of these

kernels.

Um I'll briefly highlight two works here

that um govern like that that represent

our approach. So first we think it's

important to build our own fundamental

understanding and to manually do the

work to understand it rather than just

throwing say an LLM at the problem. So

we spent the time to build out parallel

kittens which is a small set of minimal

primitives and patterns for multiGPU

kernels. Um we use this to understand

the trade-offs of multiGPU kernels and

to one write a large collection of um

peak performance kernels for a variety

of parallelism schemes. Um and this I

will use to hopefully you know educate

and bring us all on the same page on

what patterns we figured out. Um and

then once we found that there is indeed

a small set of trade-offs governing this

landscape, we were curious whether

models, especially these models right

now that claim to be very good at kernel

writing and also reasoning um could

reason about these trade-offs when we

provide them in context to actually

generate a bunch of net new multiGPU

kernels for us. Unfortunately, we found

they were not very good, but we'll dive

into more of that at the end.

Um so just the fundamental section this

is going to be more you know educational

what are the trade-offs that go into

these kernels.

So there are three main ways to do um

intraGPU data transfers. Um there's the

per GPU what's called copy engine and

this is host or CPU initiated work. It's

really good for large message transfers.

So when your message size, the amount of

data being transferred is really big um

and it can get to sort of like peak

bandwidth um on on the communication

side. In contrast, you can use device

initiated or GPU initiated transfers via

that tensor memory accelerator that I

mentioned or via register level

instructions called uh in sort of their

PTX lingo like LDST red multime. Um, and

the uh T TMA is really nice. These

device initiated ones are really nice

because they can saturate our NVLink

bandwidth using relatively small message

sizes. And this means that they can be

really nice when we're trying to do fine

grain communication rather than sending

bulk amounts of data over the links all

at once coarsely.

Um, there are some trade-offs here. So

the copy engine is really nice because

it doesn't take away or you know waste a

lot of our precious registers that I

mentioned are important for compute on

the GPU. Um and it doesn't also use any

of those rainbow colored dots the

processors on our GPU allowing us to

repurpose those for memory or um

computation on our uh you know other

parts of the AI pipeline. Um, TMA, the

second option here, consumes very few

registers, which is why it's nice. Um,

and it also can achieve high utilization

using very few of our processors. So,

it's a nice useful tool for fine grain

overlapping. Um, TMA does have

limitations. It can't effectively take

advantage of these in network um,

computations that I mentioned are

feasible with technologies like NV

switch. And the register level

instructions are really nice for being

able to take advantage of um you know

those those sort of in network

reductions that NV switch offers.

So again there are different tradeoffs

different functionalities that these

transfer mechanisms offer and they face

different trade-offs.

Second um beyond transfer mechanism the

trade-off is around how to overlap

compute memory and uh communication in

GPU kernels. So there's two main

categories of schedules. The first is

intraSM within one of those rainbow dots

um where we'll we'll have different

warps or threads within that processor

specialized to handle either compute or

one specialized for communication

concurrently. Um we can dedicate you

know different warps to each of these.

The challenge with this intm overlapping

is that the communication and

computation pattern really need to like

align and jive with one another. they

need to use the same data as inputs for

the computation and communication. When

they don't align, you could use

something like interm um schedules that

are shown on your right here where we'll

now have each of the different rainbow

colored dots on our GPUs, those

different processors specialized to

compute communication and memory. Um and

so this this is nice when the colonel

would others wise need to split across

resources like the register file or

shared memory um across these different

steps in misaligned ways.

Um this is also really nice when it's

hard to maximize NVLink traversal with

intram

overlapping.

So I just wanted to highlight one quick

example here where each of the patterns

excels in popular you know AI uh uh like

kind of patterns that you'll see. So on

a gem plus uh reduced uh scatter here we

can see that the um intraSM overlapping

scheduleuler schedule is very effective

in the gem plus all reduce we can see

that the interSM which again leverages

the in network reductions of envy switch

is very effective. So we face these

trade-offs and you can read more about

um the design decisions that go into

them in our parallel kittens paper.

Um, and then finally, ideally,

abstraction should allow the developer

flexibility to control how they're

buffering and synchronizing between data

senders and receivers.

So, we encapsulated these ideas into

parallel kittens. Again, a simple set of

programming primitives and templates for

these multiGPU kernels. Parallel Kittens

is used in production at Together AI as

well as our partner um you know Cursor

and other uh companies in the AI space.

Um here's some sample code. I won't

spend too much time here, but we usually

add roughly a dozen lines of code over a

single GPU kernel to insert these

multiGPU primitives.

Um, and you can see here across data

sequence and expert parallelism how our

parallel kittens kernels are achieving

state-of-the-art results um, compared to

strong reference baselines. And you can

check out our repo to learn more.

Okay, so we understand a little bit

about um the trade-offs that underly

these multiGPU kernels and there's just

a couple main, you know, ones that

exist. So can models reason through them

um and give us these uh you know kernels

models are getting better at reasoning

today. Um do they generalize well to

these problems or are we benchmaxed on

you know benchmarks of the past which

are more single GPU centric

uh models right now are showing really

promising results on single GPU uh

benchmarks. So it's a ripe time to to

extend it.

In our benchmark parallel kernel bench,

each task presents the model with an

unoptimized reference implementation

written in PyTorch with torch

distributed um nickel operations and

then a system topology that specifies

the number of ranks and intraode

hardware configuration. And the model

needs to rewrite the reference into a

performance CUDA kernel that uh uses

unified virtual addressing.

The um multiGPU problem space expands

combinatorally beyond single GPU um

cases. So a standard transformer layer

can be parallelized across data

sequence, tensor, context, layer,

pipeline and expert dimensions. and each

composition um induces a different

communication pass uh pattern. So to

make sure that our benchmark has high

coverage over the representative types

of multiGPU problems um we created this

taxonomy um that you can read more about

in our paper and then picked

representative problems for each part of

the taxonomy.

Um these are all patterns that arise in

real AI workloads from inference to RL

to post-training.

There are 87 problems overall um drawn

from GitHub repositories that we found

to be very informative. Um and we uh and

like optimized library implementations

and DSLs that people have written

multiGPU kernels in.

Um, we wanted to really make sure that

solving PK this parallel kernel bench

would lead to net new useful production

kernels rather than artificial or

useless kernels.

Okay, so that's parallel kernel bench.

How do models perform?

So we measured to sorry that this is a

bit small. We measured two main metrics.

Um pass at K which is the number of

correct kernels generated after K

attempts and then fast um one at K which

counts solutions that are both correct

and outperform the speed of the pi torch

plus nickel baseline.

Um so pass K just correctness fast one

at K is whether you're getting a 1x or

higher speed up over the reference. So

performanceoriented

um we found that in the zeroot setting

the best of the frontier models we tried

solves 28 out of 87 problems and 22 of

those problems are faster than the

pietorch plus nickel baseline.

If we make multiple samples, you know,

standard uh scaling up test time

compute, we can uh get that number from

say like to to 36 correct solutions, but

the fast uh one performance still

plateaus out at roughly 31%. So we don't

see much room from continuing to scale

there as we increase the number of

parallel generations.

Um we find that the correct once

correctness is established um speedups

naturally come from eliminating nickel

staging overhead in favor of direct

NVLink loads in stores.

Um the success patterns here are really

concentrated into familiar patterns. So

collective primitives, tensor parallel

gems and Ulyses style context

parallelism. So in other words, patterns

that we see heavily represented um on

the internet rather than necessarily

patterns that the model has used its

reasoning abilities to think through.

Okay, even the best available model that

we benchmarked here GPT 5.5 drops off

very quickly as the speed up threshold

increases. So on the x-axis here, we're

increasing the speed up threshold over

that pietor torch plus nickel baseline.

And then we're showing the number of

correct kernels that are faster than

that baseline or this much faster than

the baseline on the y-axis.

So GPT 5.5 is this orange line here and

then DeepSeek V4 Pro is the aqua line at

the bottom.

um we found that there's deeper issues

than CUDA syntax. So we found that if

you do multiple sampling or have the

model kind of look at its errors and

correct them, it can often compile the

kernels. But the models really struggle

to reason through the tradeoffs that we

talked about in the prior section. um

collective ordering, data partitioning,

thinking about intra versus interm

scheduling or deciding between the

different transfer mechanisms. Um we

find that they often do not use things

like the register transfer instructions

or tensor memory acceleration when

writing the kernels.

Um we wanted to try a pretty simple

instantiation of something like a clawed

code um coding agent. So, we took uh the

mini sui agent multi-turn harness and

one of the best performing models,

Gemini 3 Pro, and gave it access to a

local bash environment to sort of mimic

the standard claude code setup. We found

that this could help the agent um go

from solving 24 problems to 35 of 87

problems um with um 26 achieving over a

1x speed up over the reference. But we

found that as we scaled the amount of

time um the performance plateaued as uh

and we find that additional techniques

would be required to continue seeing the

scaling there. Again these results are

discussed in more detail in our paper.

Um we think this is a really exciting

you know just to wrap up here we hope

that people out here can use both

parallel kittens and parallel kernel

bench. We think that the kernels

generated from solving parallel kernel

bench will lead to net new production

kernels that are important bottlenecks

for inference in RL right now. um

they're you know we tried our best to

make them you know non-artificial and we

can already see signs of life and

exciting results where people have not

invested a bunch of time to handw write

a multiGPU kernel and we've gotten some

net new interesting ones like this Nemo

vocab parallel um you know filtering

kernel um a hyena architecture context

parallelism kernel

and uh the SAM 3 video segmentation

model um IOU suppression kernel.

Just to conclude here um we're really

excited about um how uh well just

talking about the lessons. First off, we

think there aren't that many patterns

that are involved in writing intragpu

effective kernels again encapsulated by

our small set of programming primitives.

Um, but unfortunately models do not

currently understand how to reason

through these trade-offs even when we

provide them in context. Um, we're

really excited about methods that can

help attack this benchmark. We're

excited about architectures that can

grow with the trends of how networking

stacks are evolving. Um, you know,

larger scale up domains shift away from

scale out um and massive onchip memory

structures. And we hope that uh you know

we can also extend and you you can feel

free to reach out to me at my email.

Thanks.