Well, there’s no easy way to say this. After about 4 years of working at PyTorch, I’ve decided to leave PyTorch to be a founding engineer at Thinking Machines. In that sentence, I would emphasize “be a founding engineer at Thinking Machines” far more than “leave PyTorch” - I have (and continue) to enjoy working on PyTorch, and would have gladly worked here another 4 years.

At several points over the last several years, I’ve talked to folks that have expressed surprise that I’m still at PyTorch. Not to brag, but it certainly wasn’t for lack of opportunity - I’ve been offered roles at OpenAI/Anthropic, I was recruited to be a founding engineer at {xAI, SSI, Adept, Inflection, etc.}, I’ve been offered many other startup roles at other startups you likely know, etc. With hindsight, many of these opportunities would have led to much greater compensation, but I’ve never regretted staying at PyTorch.

Let’s talk about why I’ve enjoyed working at PyTorch for 4 years, and what compelled me to go to Thinking Machines.

Apologies in advance for this note, it’s a self-indulgent and personal post. But I only get to write this once!

How I ended up at PyTorch

I think it’s fair to call me an AI “true believer”. Ever since I saw AlphaGo in high school and read the WaitButWhy AI post (not totally sure it holds up 10 years later), I was convinced that AI would be the most important technology of my lifetime. Correspondingly, from the time that I started college in 2016, most of what I did was AI-related. I took ML classes, I founded an undergraduate ML research club, I published papers, I even met my girlfriend (now fiancee!) through doing ML research together!

However, there were several things that didn’t totally satisfy me about just doing ML research.

For one, although I was publishing papers and such, even back then, it wasn’t super clear to me whether any research I did was actually “meaningful”. In research, one demoralizing aspect is that, looking back, 99% of papers don’t end up on the “critical path” to what actually works. Cynically, any PhD spending their time on n-gram models wasted their time - their papers and theses relegated to the dustbin of history. Although it’s true that even papers that aren’t on the critical path can still be useful and necessary (e.g. demonstrating limitations of existing approaches, setting up baselines for new approaches to surpass), this worry constantly nagged at me.

Second, I was never able to adapt well to the “running experiments” mode of ML experimentation - my working style is somewhat irregular, with lots of thinking punctuated by lots of coding. On the other hand, I think tremendous discipline is required to be a great ML experimentalist - it’s a constant loop of “come up with hypothesis” => “run experiment” => “get result of last experiment back” => “come up with new hypothesis”, often pipelined several stages deep. In ML research, you have a physical constraint (GPUs) and to be a good researcher you must get good experiment utilization out of it.

Overall, I ended up gravitating much more to “systems”. Not only was it an area that I thought played well to my strengths, I had always admired the impact of systems. Instead of needing to directly deliver the impact, you could instead improve the impact of thousands or millions of people by 5% instead!

So, I ended up with my career plan - instead of directly working on advancing ML development, I would instead work on building infra to help others advance ML development. A bunch of other things happened in-between, but that’s how I ended up on PyTorch.

Why PyTorch is an amazing place to work

PyTorch’s Impact on the Industry

As the field (and money!) has exploded over the last 10 years, I think it’s easy to lose track of just how much impact PyTorch has had. Perhaps the most obvious tracker of the money in the field is Nvidia’s stock price, primarily driven by growth in server-side GPU sales.

I think it’s reasonable to guess that at least 75% of these GPUs are running some kind of PyTorch code.1 That’s insane. Nvidia has gained something like 3T dollars of market cap, and PyTorch was crucial for much of it.

Moreover, among the ML community broadly, PyTorch remains the lingua-franca. 59% of research papers tracked by Papers With Code use PyTorch (with 29% not using any ML framework), the vast majority of models on Huggingface (90%+?) are on top of PyTorch, and the most popular inference servers are also built on top of PyTorch (vllm and sglang).

Even at the leading AI labs, pretty much all of the companies using GPUs use PyTorch. OpenAI, Mistral, Deepseek, and Meta primarily use PyTorch (and GPUs). Anthropic also primarily uses PyTorch for GPUs, and xAI (which uses Jax for training on GPUs) also uses PyTorch for inference (through sglang)!

In high school, one of the things I feared most was that I would work on some project for 10 years and eventually realize that I’ve wasted my life improving something that nobody cared about. One of the greatest things about working on PyTorch is the certainty that I haven’t.

PyTorch’s Impact on Me

I’ve spent my entire career (thus far) on PyTorch, and so, outside of the overall impact of PyTorch, I want to talk about why I’ve enjoyed the day-to-day so much.

Mission Alignment

One of the greatest things about startups is “mission alignment”. Because so much of your comp is tied to equity upside, there’s no difference between “my coworker wildly succeeded” and “we all wildly succeeded”. On the other hand, at a large tech company, people’s compensation is primarily tied to their individual performance rating (and promos). So, if you start working on an approach, and somebody else comes up with a different approach that’s wildly successful (and supersedes yours), your performance rating is likely to suffer, and you probably won’t get promoted.

At PyTorch however, many of the people on the project are mission aligned - they do genuinely care about the overall success of PyTorch and its impact on the ML ecosystem. I certainly wouldn’t say it’s 100% of people on the team, but it’s enough (especially among more senior folk) to make it a much more pleasant experience.

A True Commitment to Open-Source

Soumith (and other folks in leadership) have done an exceptional job in cultivating a culture where OSS is valued at Pytorch. There are many other projects that happen to be OSS, but you can often only get promoted and have impact through prioritizing internal projects. This isn’t true at PyTorch - I would say that I’ve spent my entire time here primarily focused on OSS impact, and I’ve been successful in terms of ratings and promotions. (Of course, there are other folks primarily focused on internal impact who have also been very successful).

There are other aspects in which valuing OSS leads to a much healthier project.

Ungameable Impact

One phenomenon at large tech companies that I don’t really like is what I call “roadmap-driven adoption”. This is where two managers/directors/VPs get together, agree that X should be used (potentially killing other project Y), and then the adoption of a project is laid out in roadmaps of several teams.

While this certainly has its advantages (and in some cases is entirely necessary), I find that projects adopted this way are often … subpar. Moreover, it’s not uncommon for the success of these projects to be a facade - they continue as long as some VP is sponsoring the project, but eventually people get fed up with it, the VP loses a political fight, or the VP simply changes their mind. Basically, in roadmap-driven development, the most important component is convincing some “key stakeholders” that your project should be adopted.

On the other hand, OSS impact is truly a free market. Open source users couldn’t care less if Mark Zuckerberg is throwing his entire support behind a project. OSS users only care that 1. you’re solving a problem they have, and 2. They like using your software.

Mike Schroepfer (former CTO at Meta) provided a similar sentiment here. I can’t even imagine how hard it is to get “real” feedback as a CTO, where everybody you talk to knows that you can be single-handedly responsible for their promotion or bonus. OSS is refreshingly ungameable.

Amazing Opportunities for Career Growth

Even from a purely selfish perspective, it’s hard to imagine a place with more opportunities for career growth. Folks can often get tunnel vision on internal promotions. But, a much bigger driver of career growth is often “what do folks across the industry think of you”. Objectively speaking, I’ve gotten incredible offers as founding engineers at startups (SSI, xAI, Thinking Machines) as well as great opportunities for (fairly) high profile talks (Jane Street, a couple lectures at universities, etc.). Even among big tech companies, at some point a director offered me a 2.5 level jump to go work at Google.

Frankly, it’s embarrassing to write the above paragraph - there are definitely other folks on PyTorch who are just as good or better than me, but due to my focus on OSS and public presence (twitter/blogs/etc.), I’ve become fairly well known. But… this is somewhat of a “sell-post” on PyTorch, so I’ll allow it :)

I would say that just working on PyTorch alone is a great opportunity. Andy Jones (at Anthropic) once wrote an essay whose top line was “If you think you could write a substantial pull request for a major machine learning library, then major AI safety labs want to interview you today.” Obviously, if you work on PyTorch, you qualify for this.

But where Pytorch truly excels is the opportunity for engineers to have industry-wide impact (and to get recognized for it!). It’s very rare that most people in the most important field of our time use your software - I’d recommend taking advantage of it.

Last note - even for internal work, I find broad OSS impact is very helpful. Especially when it comes to cross-team collaboration, one of the most important currencies is “legitimacy”, and OSS impact is far more likely to confer legitimacy than internal impact. I’ve benefited significantly from this.

Interesting Technical Work

One fear of many engineers is that they won’t be able to solve interesting technical problems - there’s no shortage of that on PyTorch. There’s projects that have implemented a python bytecode interpreter JIT for machine learning (TorchDynamo), projects on reaching speed-of-light for matrix multiplications, projects where you regularly need to dive into PTX documentation (perhaps this one isn’t appealing hahaha), projects all about reasoning over symbolic shapes (sympy, z3, etc.), and so many more.

There’s also no shortage of problems to be solved 😛

Consider Working on PyTorch!

Perhaps a year ago I was talking with Adam Paszke about working on ML frameworks, and we agreed that it was surprising to us that given how sweet the gig is, it was surprising to us that more people didn’t want to work on ML frameworks. This post is (partially) my attempt to remedy that.

If any of the above sounds appealing to you, consider working on the Pytorch Core2team (primarily under Gregory Chanan)! I would recommend contacting/emailing Soumith Chintala at soumith@meta.com. Although I will (sadly) no longer be on the team, there are many other amazing people to work with (and I plan to continue to be in regular contact with the team).

From my perspective, an ideal candidate would:

1. Have an inherent curiosity in the field of machine learning: Machine learning moves very quickly, and one characteristic I’ve found most valuable is a general knowledge of what’s going to be useful next. For me, I found the years of just doing pure ML research (using PyTorch) to be extremely helpful in understanding what should be built.

2. Be highly agentic: Although having ideas on what should be done is an important characteristic, it’s perhaps even more valuable to just “do it”. PyTorch, in many ways, is a very “bottom-up” org. Although there’s always roadmapping and such, the most impactful projects were never on the roadmap.

Why I’m Excited About Thinking Machines

Given that I just wrote way too much about why I loved working on PyTorch, why did I join Thinking Machines? Moreover, why was Thinking Machines the opportunity that convinced me?

A Group of People I Would Very Much Like to Work With

As everyone knows, a startup is nothing without its people. And Thinking Machines sure has some pretty good people!

You have the folks responsible for the “research preview” that’s kickstarted this situation (John Schulman, Barrett Zoph, Luke Metz), you have folks leading pretraining efforts at Meta, OpenAI, Character.AI, etc (Sam Schoenholz, Naman Goyal, Myle Ott, Jacob Menick), you have folks leading multimodal efforts at OpenAI/Mistral (Alexander Kirillov, Rowan Zellers, Devendra Chaplot), you have extremely good infra people (Andrew Tulloch, Yinghai Lu, Ian O’Connell, etc.), and of course you have the former CTO (and brief CEO :^)) of the biggest AI company in the world (Mira Murati).3

However, perhaps even more so than the strength of the team, I was impressed by the friendliness of the team. Certainly, the fact that I had previously enjoyed working with several of these people helped. Of the top 4 people who I’ve been saddest about leaving Meta, 2 of them (Andrew Tulloch4 and Yinghai Lu) are at Thinking Machines.5

An Amazing (and Asymmetrical) Opportunity

One unfair advantage about being a founding engineer at a startup (especially one that’s so clearly a “good” opportunity) is the asymmetrical opportunity cost. 6

For example, if I join Thinking Machines as a founding engineer, and then in a year decide I was massively mistaken and go to another lab, it’s not likely that my role would change all that much! I would be joining an established company, and the role would likely be fairly similar to what it is today. Somewhat macabrely, just like the prodigal son, it might even be beneficial to have left.

However, if I declined now but then joined Thinking Machines in a year, my role would be drastically different. Of course, my compensation would change, but more importantly, I would have far less legitimacy and influence. The culture and direction of a company is largely set by the founding team, and that’s something I don’t have the opportunity for at OpenAI or Anthropic.

An Approach to Positive AI Outcomes that Resonated with Me

Perhaps most importantly, however, was that Thinking Machines’ approach to positive AI outcomes - research/product codesign and open science - resonated with me. As mentioned above, I’ve been convinced since high school that AI would be the most important technology of our lifetime. However, this is not the same thing as saying it’ll be the most beneficial.

Note: Don’t take the below statements as speaking on behalf of Thinking Machines - although I’ve talked to members of the team about these various topics, this is entirely “my stance” and “why does my stance lead to me joining Thinking Machines”. I also don’t have the space to give this topic a more in-depth treatment, so view it more as an “explanation of my beliefs” as opposed to a watertight argument.

Generally speaking, I would consider myself a techno-optimist. That is, I believe that humans’ lives have gotten drastically better over the last 1000 years, and that this has been largely driven by technological innovation. I’m thankful for the modern supermarket, I’m astounded by semiconductor manufacturing, and I’m especially grateful for air conditioning. In many ways, AI is the most techno-accelerationist technology the world has ever seen - a single technology that has the potential to solve every other technical challenge we face. Because of this, the potential positive impacts of AI are worth pursuing - I like Dario’s piece “Machines of Loving Grace” for a more in-depth look on what that might look like.

Of course, bad outcomes are possible, and due to the potential impact of AI, bad outcomes seem far worse than with other technologies.

In general, I’ve categorized bad AI outcomes as:

• Misuse: Bad people use AI to do something bad7

• Misalignment: Good people use AI, but the AI itself ends up doing something bad.8

• Societal Impacts: People are good, AI is good, but we end up in a bad outcome anyways.9

Although all 3 are reasonable concerns, I would say I’m most concerned with “Societal Impacts”.10

The primary reason for this is that society naturally has a strong “immune response” to Misalignment and Misuse. When it comes to potentially harmful technologies, society has a clear playbook - if something bad happens, increase the restrictions (e.g. regulate GPUs) or regulations (e.g. mandate more safety oversight).

Of course, AI is not a normal technology, but concretely speaking, I think there will be plenty of warning signs before truly catastrophic misuse or misalignment. Even if the AI bides its time before misalignment (e.g. deceptive misalignment or a treacherous turn), I find it unlikely that the first AI system to do so will succeed - it would need to be drastically more powerful than humans and other AIs.

On the other hand, negative Societal Impacts seems much more straightforwardly plausible. Imagine a world where ChatGPT was never released, but instead hypothetical AI lab ClosedAI releases AI Job Replacement Agent 3000 in 2030, instantly replacing 50% of all human jobs and turning ClosedAI into a quadrillion dollar company.

Even in today’s world, the secrecy of the top AI labs definitely rubs me the wrong way (although I understand why it’s done) - I don’t know how much more vagueposting I can take. Moreover, the ideological and geographical concentration of AI knowledge doesn’t seem ideal - as AI expertise becomes more and more in-demand, the fact that the vast majority of AI secrets are contained in a 50 mile radius around San Francisco leads to both power-imbalance and a monoculture.

If we need to align AI to human values, should all of those humans live in San Francisco?

Why I’m Compelled by Thinking Machines’ Mission

Broadly speaking, there are 2 main aspects of Thinking Machines’ mission that were compelling to me.

A Focus on Product and Broad AI Diffusion

In my opinion, one of the most important aspects for broader societal stability is how smoothly society transitions to using AI systems. Just as important as the outcome is how people feel we arrived at that outcome.

For example, ChatGPT did not really blow many ML researchers’ minds - they’d seen GPT-3, they’d seen what GPT-3 prompting could do, ChatGPT was just a convenience feature. However, ChatGPT absolutely blew the rest of society away. This was the first time broader society became aware of all the things a SOTA LLM could do, and society was shocked. However, since, ChatGPT has become much more normalized among broader society - folks have a bit of a hedonic treadmill.

But, there is much more that can be done. Even today, there is a vast gap between what the layman encountering ChatGPT for the first time can do vs. those who have deeply integrated AI into their workflow.

Moreover, I believe there is a lot of potential towards building AI products that can help out people collaboratively as opposed to fully autonomous AI agents. One cute way I thought about it was “Maximize the value of labor instead of capital”.

Open Science and Systems

As mentioned above, it doesn’t seem good for society for the knowledge of how these AI systems are built to be so secretive. Not only does it create resentment towards these AI labs, it also makes it much more difficult to society to build on top of these AI systems! For example, Deepseek’s recent releases (both their papers and code) have helped the broader community develop a much better understanding of what will be useful moving forward (e.g. Online RL).

Personally, of course, this was a big part of my motivation for PyTorch. Good open-source systems help out the entire ecosystem, enabling many more people to participate in building AI systems.

I also want to note that although open science/systems is certainly a nice ideal, there’s obviously also economic realities at play. In my opinion, this is where the focus on product is useful. Companies like Meta or Google don’t need to be very secretive about the actual techniques being used - more or less, most of their core systems/approaches are widely known by the community. On the other hand, if your only product is just an API endpoint with tokens in and tokens out, your only edge is your model’s exact capabilities.

The culture and defaults of a company also matter a ton. There are many things at these AI labs that could be OSS’ed without affecting their competitive edge - they just don’t do it because the default is closed and they’d need to argue why it should be open.

For example, PyTorch is the opposite here. All of our code is open-source, our roadmaps are open, some of our design meetings are open too. So, one must argue why something should be closed if you don’t want it to be open.

Based off of Sam Altman’s comment here, he thinks that OpenAI should be open-sourcing more things. However, it’s “not [the] current highest priority”.

Overall Thoughts on Positive AI Outcomes

Overall, I think Thinking Machines mission of broad AI diffusion and collaborative open-science seems like a compelling strategy to help address Societal Impacts. Of course, there are other essential approaches (like Policy), but Thinking Machines’ mission personally resonates and is an area I believe I can contribute to.

Final Thoughts

The opportunity to join Thinking Machines as a founding engineer hit basically all of my checkboxes.

1. An extremely strong team, with folks I’ve personally enjoyed working with before and other folks I think I’ll enjoy working with.

2. The opportunity to be there from the beginning and have a say in the direction and culture of a very exciting company.

3. A mission (product-focus + open science) that was uniquely compelling to me in leading to better AI outcomes.

4. Finally, as a bit of a gut-instinct thing, the open science/systems aspect allows me to continue some of what I enjoy about my role at PyTorch - talking to people about AI systems and having broad impact with open-source code.

Almost none of my previous opportunities hit even 2 of these boxes, let alone all 4. One point I distinctly remember when considering this was: “If even this opportunity doesn’t compel me to leave PyTorch, I should/would probably work on PyTorch forever”.

Although it was a very difficult decision, I’m very excited to build some cool stuff at Thinking Machines!

Come work at Thinking Machines!

If any of the above sounds compelling to you (or if you want to work with me :P), please email me at horace@thinkingmachines.ai or DM me on twitter!

At the very least, I think it’ll be a lot of fun!

Acknowledgements

I’d like to thank everyone I talked to while making this decision, all of which provided me a lot of insight. I’d also like to thank my coworkers at PyTorch for making my last 4 years so enjoyable.

The beginning excerpted here:

In theory, Attention is All You Need. In practice, however, we also need optimized attention implementations like FlashAttention.

Although these fused attention implementations have substantially improved performance and enabled long contexts, this efficiency has come with a loss of flexibility. You can no longer try out a new attention variant by writing a few PyTorch operators - you often need to write a new custom kernel! This operates as a sort of “software lottery” for ML researchers - if your attention variant doesn’t fit into one of the existing optimized kernels, you’re doomed to slow runtime and CUDA OOMs.

For some examples of attention variants, we have Causal, Relative Positional Embeddings, Alibi, Sliding Window Attention, PrefixLM, Document Masking/Sample Packing/Jagged Tensors, Tanh Soft-Capping, PagedAttention, etc. Even worse, folks often want combinations of these! Sliding Window Attention + Document Masking + Causal + Context Parallelism? Or what about PagedAttention + Sliding Window + Tanh Soft-Capping?

The left picture below represents the state of the world today - some combinations of masking + biases + setting have existing kernels implemented. But the various options lead to an exponential number of settings, and so overall we end up with fairly spotty support. Even worse, new attention variants researchers come up with will have zero support.

To solve this hypercube problem once and for all, we introduce FlexAttention, a new PyTorch API.

1. We provide a flexible API that allows implementing many attention variants (including all the ones mentioned in the blog post so far) in a few lines of idiomatic PyTorch code.

2. We lower this into a fused FlashAttention kernel through torch.compile, generating a FlashAttention kernel that doesn’t materialize any extra memory and has performance competitive with handwritten ones.

3. We also automatically generate the backwards pass, leveraging PyTorch’s autograd machinery.

4. Finally, we can also take advantage of sparsity in the attention mask, resulting in significant improvements over standard attention implementations.

With FlexAttention, we hope that trying new attention variants will only be limited by your imagination.

My next post (~70% done) will also be about attention! This one will be a historical retrospective on why Tri Dao was the one to invent FlashAttention, and not any of the large tech companies.

There’s some other interesting topics I’m considering writing about.

In particular, some potential titles:

1. What’s the point of ML compilers when Attention is All You Need? (how I think about building ML systems in a world where everybody uses one ML architecture)

2. Performance Metrics Were Made for Man, not Man for Performance Metrics (How do you choose the right performance metric?. In particular, flop counting)

3. My ML framework isn’t Obeying Mathematics! (a primer on floating point “nondeterminism” for machine learning settings)

python mm_bench.py
> CuBLAS: 258 Teraflops

Not bad, 83% flop utilization. Now let’s check out Cutlass’s performance using their profiler.

./cutlass_profiler --operation=Gemm --m=8192 --n=8192 --k=8192
> CUTLASS: 288  Teraflops

!!! 10% higher perf? That’s incredible. CuBLAS is highly optimized for large compute-bound matmuls, and somehow CUTLASS + autotuning is outperforming it by 10%? We gotta start using these matmuls yesterday.

The next step is to bind the CUTLASS kernels into Python and compare against CuBLAS using my previous script.

python cutlass_mm_bench.py
> CuBLAS: 258 Teraflops
> CUTLASS: 257 Teraflops

Somehow, in the light of Python, all of CUTLASS’s performance gains disappear. This in of itself is not shocking - it’s notoriously difficult to ensure consistent benchmarking across setups.

I tediously ablate the two benchmark scripts, until finally, I find that CUTLASS’s profiler, by default, actually initializes the values in a fairly strange way - it only initializes the inputs with integers. Confused about whether this matters, I try:

zero_inputs = torch.zeros(N, N)
randn_inputs = torch.randn(N, N)
benchmark(zero_inputs) # 295 Teraflops
benchmark(randn_inputs) # 257 Teraflops

What? How could the values of the matrix affect the runtime of the model? I know Nvidia has some weird data compression thing on A100s, but I wouldn’t have expected that to be on in matmuls. Let’s try some other data distributions, like an uniform distribution [0,1].

This was … confusing, to say the least. Somehow, the actual content of the tensors being multiplied is leading to different matmul performance.

There certainly are cases where the runtime depends on the content of the tensor — indirect indexing (e.g. A[b]), or things like sparsity.

But matrix multiplications have nothing like that at all! No matter what the contents of the matrix contain, the matrix multiplication kernel will 1. perform the same number of computations, 2. perform the same computations in the same order, 3. access the same memory addresses, and 4. access the same memory addresses in the same order.

Nowhere did my mental model of matrix multiplications and GPU hardware allow for the values in the matrix to influence matmul performance. And yet, here we are.

As it turns out, the culprit is ……. dynamic/switching power in semiconductors!

Power Usage in Semiconductors

An Nvidia A100 GPU has a power limit of 400W1. However, as the phrase “power limit” may hint, the GPU doesn’t always use all 400W. For example, when the GPU is fully idle, nvidia-smi tells me that it’s only pulling 88W of power.

But when the GPU is running under load, that power usage will spike considerably, typically to around the power limit.

In order to stay under the power limit, a piece on the chip called the Voltage Regulator Module reduces the voltage supplied to the GPU, — throttling the clock frequency and reducing its performance.

In other words, if our GPU ends up using enough power to hit the power limit, our performance will become capped.

Most of us take it for granted that “GPU does something, power consumption goes up”. But there are actually two distinct mechanisms through which power gets consumed.

The first one is static/leakage power. You can think of this as the power that inevitably gets lost by just flowing power through your circuits. The amount of static power used is proportional the amount of silicon that is powered. As GPUs don’t do much power gating, this is essentially the amount of power used at idle (88W in the above photo).

However, the second one, dynamic (or switching) power, is the culprit. Specifically, a small amount of power is consumed whenever a transistor switches states. If the transistor never needs to switch states, it doesn’t consume any extra power. On the other hand, if it’s rapidly flipping, then it consumes a ton of dynamic/switching power. Multiply that by the billions of transistors in your GPU, and you get the overall increase in power consumption.

In other words, the reason why matrix multiplications are faster when passed zeros is that this reduces the “flipping” of enough transistors in the chip to stay under the power limit!

So, this (mostly) explains what we saw previously2. All zeros are probably the fastest since every single bit of each computation is a zero and the accumulator remains at zero. All ones is probably still quite fast since every single tensor-core instruction results in exactly the same values. The uniform distribution probably is a little bit faster than the normal distribution since the accumulator never needs to flip-flop between positive and negative. The normal distribution probably has the worst performance since it leads to pretty high randomness among all transistors involved in the computation(?).

Here’s the results on a number of fun distributions I tried:

1. Randn: Normally distributed

2. Checkerboard: Normal distribution, but we have zeros in a checkerboard shape.

3. Rand: Uniform distribution

4. Sparse: Normal distribution, but a (random) 75% of the elements are masked.

5. Ternary: Every value is 1, -1, or 0.

6. One Bit: Only one bit set in every single value (the 4th bit)

7. All Pies: Every single value is the mathematical constant PI.

8. Twos: Every value in the matrices are 2

9. Zeros: Every value in the matrices are 0

Who says unstructured sparsity isn’t efficient with tensor-cores? :)

How power limit and clock speed affects this

Here’s another piece of evidence that dynamic/switching power is responsible.

Roughly, the power we’re using is proportional to the clock speed multiplied by amount of transistor flips we’re doing.

power ~= clock speed * “transistor flips per clock”.

We run into throttling when the power we use surpasses the power limit we’ve set. Thus:

1. If we reduce our power limit we exacerbate this effect.

2. If we reduce the clock speed we reduce this effect.

Let’s show that in action! To do so, I’ll compare the relative performance of a very predictable input (zeros) vs a very unpredictable input (randn).

As expected, we see that as the power limit decreases from 330W down to 100W (the minimum), the relevant performance improvement from using predictable inputs increases. Interestingly, at the lowest power limit (100W), the trend reverses. I’m guessing that the GPU is so power constrained that even using all zeros still results in too much power usage. [Uma Kelkar in the comments suggests that this is because throttling the power limit also just directly throttles the clock speeds]. Remember that the switching power is coming from every transistor in the GPU, not just the ones holding data! So that includes the transistors for say, keeping track of the current program counter, keeping track of how many loop iterations you need to perform, the ones signaling other transistors to perform operations, pretty much everything that a GPU can possibly be doing.

Now, to test the effect of GPU clocks on using predictable vs. unpredictable inputs, I’ll use a a power limit of 200 and vary the GPU clock limit.

We see that at the top range of the clock frequency, there is nearly no change in the ratio, as presumably even with predictable inputs, we’re still getting throttled. Then, as we decrease the clock speed, the relative gap shrinks as predictable inputs are affected by the clock speed limit, but not unpredictable inputs. Finally, at the very left of the chart, both using predictable and unpredictable inputs have identical performance, as both become completely limited by our manual clock speed limit and don’t do any power throttling.

Another interesting thing we can test is, for a given input and power limit, what is the maximum clock speed the GPU can sustain for a matmul?

Marketing vs “Real” Performance

This observation that GPUs are unable to sustain their peak clock speed due to power throttling is one of the primary factors that separates “real” matmul performance from Nvidia’s marketed specs.

The figure that Nvidia provides for marketing is:

For example, on an H100, there are 528 tensor cores per GPU (4 per SM), the max clock speed for these is 1.830 Ghz, and the FLOP per tensor-core instruction is 1024. Thus, we have 1.830e9 * 528 * 1024 = 989 TFLOPS, exactly Nvidia’s listed number.

However, you can only achieve this number by sustaining 1.83 Ghz clocks, and as we’ve seen above, the GPU just doesn’t have enough power to do that!

Do note that in both of these cases, the GPUs have a higher power limit than I can test (400W on A100 and 700W on H100 respectively), so it’s able to sustain a higher clock speed than what’s charted here. But we can see that, especially on the H100, the max sustainable clock speed is much lower than the theoretical one! In other words, matmuls on the H100 are primarily not compute or bandwidth limited, they are power limited.

As many have noted, power is increasingly a crucial constraint. So, although the H100 theoretically had 3x more FLOPS than the A100, its “real” performance has usually been closer to 2x due to the power throttling we’ve been discussing, and its “flops per watt” is even less than that.

Update (6/25): Nvidia recently released their MLPerf submission which largely confirms the hypothesis in this blog post. Thanks Sophia Wisdom for the pointer!

This leads to high Tensor Core utilization and can result in scenarios where Tensor Core throughput is constrained by the power available to the GPU. 

In the submission with 512 H100 GPUs, we improved end-to-end performance by redirecting power from the L2 cache memory on each H100 GPU to the streaming multiprocessor (SM), which houses, among other units, NVIDIA Hopper fourth-generation Tensor Cores.

Conclusion

All of this should make you exceedingly curious to see the actual performance improvement on the B100, which has a 1.75x theoretical increase in FLOPS with the same power usage as the H100. I have a Manifold market about guessing the max FLOPS utilization on a B100.

I’ll leave you with a slightly modified version of this tweet from roon.

Thanks to Sophia Wisdom, Vijay Thakkar, Philippe Tillet, Dylan Patel, and Natalia Gimelshein who have helped me understand this phenomenon.

Also, be aware that you need to set the `scale` parameter on the CUTLASS profiler if you wanna compare its FLOPS numbers to other benchmarks!

Question 1

Let's say I have a [M x K] @ [K x N] matmul. Which one of these configurations will have the best perf? Think about the actual ramifications of tiling! Both matrices are in row-major layout (i.e. K and N are the innermost dimensions)
A: M=2047, K=N=2048
B: K=2047, M=N=2048
C: N=2047, M=K=2048

Answer

The correct answer is A. The key to this question is understanding the point about tiling and memory layouts.

When it comes to memory layouts, not all dimensions are created equal. Matrix multiplications are not inherently allergic to odd numbers - the poor performance is for a very specific reason.

In this case, that specific reason, as described here, is that odd shapes lead to unaligned memory layouts. Crucially, however, it must be an odd innermost dimension that leads to unaligned memory layouts.

Before explaining it, let’s have a practical demonstration that A (i.e. M=2047, K=N=2048) is over 2x faster than either alternative. Code can be found here.

import torch
from triton.testing import do_bench
torch.set_default_device('cuda')

for M, K, N in [(2047, 2048, 2048), (2048, 2047, 2048), (2048, 2048, 2047)]:
    A = torch.randn(M, K, dtype=torch.bfloat16)
    B = torch.randn(K, N, dtype=torch.bfloat16)
    print(f"M={M}, K={K}, N={N}")
    print(do_bench(lambda: torch.mm(A, B)))

To understand why this occurs, let’s understand how the logical layout and the physical layout look like.

First off, let’s say we have an 8x8 matrix, nicely aligned. Pretend that each cache line is 4 elements long.

This is our ideal situation. The physical layout lines up with the cache lines, each load perfectly uses up every element in our cache line, and the world is at harmony.

However, let’s say we introduce an extra element per row, resulting in an 8x9 matrix.

This one measly element per row throws everything out of balance. Each row no longer starts on a cache line, and when issuing our loads, we can no longer simply load from a single cache line to obtain all the elements we need. (This is all just a restatement of the explanation from here)

However, what happens if instead of adding an extra element per row, we simply add an extra row, resulting in a 9x8 matrix?

Unlike before, this does not affect the “alignedness” of each row! We do have an extra row at the bottom, and this may lead to some extra computation, but if our matrix was sufficiently large, that computation would be negligible. The important point, however, is that this extra row does not affect the memory layout of the rest of the matrix.

In other words, as long as the innermost size of your matrix is divisible by the cache line size, you’re good to go!

So, armed with our refined understanding of how shapes affect matrix multiplication performance (e.g. only the “evenness” of the innermost dimension matters for memory layouts), let’s look at the question again.

In a matmul, we have A: [M x K] and B: [K x N]. These are both row-major, which means that for K and N are the innermost dimensions of A and B respectively.

A: M=2047, K=N=2048 (the right answer!)
B: K=2047, M=N=2048 (Ruled out because K is the innermost dimension of A)
C: N=2047, M=K=2048 (Ruled out because N is the innermost dimension of B)

Interestingly, when I originally asked this on Twitter, most people got it wrong.

The substack readers fared much better.

One friend answered that they chose B because “A and C seemed symmetrical, and so B must be the right option, since A and C couldn’t both be right”. Sadly, especially in the world of systems, things that seem identical may not be identical in practice…

I added Question 4 because some people were getting the right answer for Question 1 but for the wrong reasons, so let’s jump ahead and see how our newfound knowledge applies to a slightly modified version.

Question 4

Similar to Question 1, let’s say we have a A: [M x K] @ B: [K x N] matmul. However, now, A is in column-major (i.e. torch.randn(K, M).t()) while B is still row-major. What is the best configuration now?

A: M=2047, K=N=2048
B: K=2047, M=N=2048
C: N=2047, M=K=2048

Answer

Read more

This may seem intuitively quite strange. Is cuBLAS just messing up somehow? Why doesn’t the matrix multiplication kernel just pad it to a larger shape?

It has become tribal knowledge that the particular shapes chosen for matmuls has a surprisingly large effect on their performance. But … why? Can this be understood by mere mortals?

Let’s take a crack at it.

First, let’s plot FLOPs achieved for square matmuls. By the end of this article, I will aim to explain all the strange squiggly lines.

There are 3 general concepts to understand that explain the majority of performance variation among matmul shapes.

1. Compute Intensity/Parallelization: This explains the general upward trend

2. Tiling: This explains the multiple tiers of lines.

3. Wave Quantization: This explains the strange striped lines.

Compute Intensity and More Parallelism

First of all, as we move along the x-axis, the matrix multiplications generally get more performant. There’s two primary reasons for this.

The first one is simply “more work/more parallelism”. There are a large number of fixed overheads that come with launching a kernel (e.g. creating new SMs, waiting for all SMs to finish, etc.), and so, the more work we have to do, the less important those fixed overheads are. Along with more work comes more parallelism, and since GPUs have a ton of parallel cores, you need a surprising amount of work in order to fill a GPU up with enough parallelism.

The second one is “arithmetic intensity”. As I’ve written about before, memory accesses are much more expensive than compute. So, since a square matmul performs 3N^2 memory accesses and 2N^3 FLOPs, at a very minimum, N needs to be in the hundreds before we start spending more time on compute than memory!

The desire for sufficient Arithmetic Intensity and Parallelism also compound. For example, let’s say you have your output matrix is 1024 x 1024. If you let each SM compute a 128 x 128 slice of the output, that’s only 64 pieces of “work” for your GPU, not even enough for each one of an A100’s 108 SMs ! If you decrease your output slice size to 64 x 64, we now have 256 pieces of “work” for our GPU, but our arithmetic intensity has also decreased by a factor of 2.

With smaller matrix sizes, you need to worry about problems like this that don’t show up with larger matrices.

Tiling

Now that we understand the overall structure of the plot, the next question is: why is the plot all over the place? Why, even for very large matrices, do the TFLOPS jumping between >250 and <100?

To give a hint, let’s color-code each dot by the highest power of 2 it’s divisible by.

As it turns out, the multiple “levels” of FLOPS are due to their shapes’ divisibility. For example, when the shape is odd, the matmul performs significantly worse than when the shape is even. The matmul performs even better when the shape is divisible by 8, with even more performance gains when it’s divisible by 16 or 32.

Now, merely knowing about this effect is very practically useful, but what actually causes this effect? As it turns out, the answer is tiling. But, what even is tiling? And why does it cause such substantial performance issues?

Some online have mentioned tile quantization as the culprit. Tile quantization certainly can impact performance, but only at tile boundary sizes. Basically, tile quantization occurs when the size of your matrix multiplication increases such that the GPU needs to launch another “chunk” of work. For example, imagine that you could multiply 8 elements at a time with a SIMD instruction. Now, if you went from 32 elements to 33 elements (a 3% increase in problem size), you go from needing 4 SIMD instructions to 5 (a 25% increase). Note that crucially, when tile quantization is the culprit, your absolute runtime still grows monotonically, although your efficiency may drop.

However, in our above plot, we see much more drastic performance drops! Moreover, like in Karpathy’s original example, we see that the absolute runtime decreases despite problem size increasing. So, tile quantization cannot be the explanation here.

The true cause is that tiling is just fundamentally worse for certain memory layouts. In other words, by the time we’re trying to execute the matmul, you’ve already lost. The memory layout is poor and your performance will suffer.

Let’s look at some examples!

Memory Layout of Tiling

First, let’s think about how our matrix’s memory layout looks like when our size is a multiple of the cache line (pretend it’s 4 elements).

We see that each row starts on a cache line1. Among other advantages, this means that we don’t need to perform any “unnecessary” loads to load all yellow elements. We can just load the 3 cache lines that the yellow elements are part of.

However, what happens if we increase the number of elements per row from 12 to 13?

With an unaligned layout, each row is misaligned relative to our cache line. In other words, if we start loading the beginning of the green row, we must redundantly load the last element of the blue row as well.

Now, let’s look at what happens when we actually try to load an entire “tile” from these memory layouts.

With the aligned layout, this is very clean! We issue one load per row. One for the 4 blue elements, one for the 4 green elements, one for the 4 yellow elements, and one for the 4 pink elements.

With the unaligned layout, things are much messier. For example, in order to load the first 4 green elements, we must issue 2 loads! One that gets the last blue element + the first 3 green elements, and one that gets the 4th green element. A similar pattern occurs with loading the 4 yellow elements as well as the 4 pink elements.

So, when our matrix size is divisible by the cache line (which is 32 elements on a GPU), tiling fits nicely within the cache line, and our memory loads are maximally efficient. When it’s not… the kernel needs many more workarounds in order to end up the proper alignment.2

This is why even very small changes in our matrix size can lead to substantially worsened performance.

Wave Quantization

Ok, so we’ve understood most of the variation in matmul performance. But what about these strange stripes up here? All of these points are with matmuls that are divisible by 32 already. Seeing that the peaks are separated by 256, our first guess might be that this is also memory-layout related, just at a larger scale.

However, as it turns out, these peaks (2944 and 3120) do not occur when the matrix shapes are divisible by 256, but instead they’re at 128 mod 256!

As it turns out, these peaks are not caused by poor memory-layouts, they’re instead caused by a (neatly-named) phenomenon called wave quantization.

The main idea behind wave quantization is quite simple.

Let’s say we have N parallel tasks (which each take a second) and N CPUs.
Q: How long does it perform to take all tasks?
A: 1 second

Q: What about if we have (N+1) parallel tasks, and N CPUs?
A: 2 seconds(!) Now, one CPU must perform two tasks, taking a total of 2 seconds.

So, despite adding just one task, we’ve doubled our overall latency.

This is exactly what wave quantization is, except with CPUs => SMs and tasks => thread blocks.

As your matrix size increases, the total number of tiles/blocks increases. When this crosses a multiple of the # of SMs, your perf drops since you need to execute an additional "wave".

Now, let's apply our newfound knowledge to actually explain these curves! Let’s try looking at this sudden drop in performance around 1792 first.

Since wave quantization depends a lot on the actual kernel parameters, we must look at what kernels are actually being run.

Using the profiler, we see that we're running a CUTLASS-based matmul with a tile size of 256x128. Note that our matmul kernel *doesn't change at all*, but our perf drops from 60+ TF/s at N=1791 to 43 TF/s at N=1793.

Now, some basic arithmetic. Our tile grid has dimensions 1792/256 = 7 and 1792/128 = 14. That gives us 7 * 14 = 98 tiles. Since an A100 has 108 SMs, that's still one wave. However, with N=1793 we need to increase the size of our grid. (7+1)*(14+1) = 120 tiles, or 2 waves!

Now, let’s look at the previous (mysterious) stripes. Specifically, we’ll look at N=3200.

Profiling it, we see that the proximal cause is not actually wave quantization. Instead, CuBLAS decided to change algorithms. But, why did CuBLAS decide to change algorithms?

Well, (3200/128) * (3200/128) = 625. 625/108 = 5.8 waves. Thus, at N=3232 we would create another wave.

In this case, though, it seems that 160x128 still isn't a great tile size. Since the resulting grid (26x21) results in 5.05 waves...

Well, CuBLAS isn't perfect!

Beyond the obvious matrix multiplication shape issues, performance loss due to wave quantization often ends up being tricky to find, since it depends upon things like the batch size as well. However, if you take a closer look at each matmul, you might find that there’s another 10-15% performance you can squeeze out of it by choosing the shapes more carefully!

I will note that it’s possible that wave quantization effects may soon be a thing of the past. New matrix multiplication technology like stream-k allow us to completely bypass wave quantization effects. Perhaps I’ll explain the basic idea behind matmul implementation strategies someday.

Why doesn’t torch.compile just fix my problems so I don’t have to think about this?

As it turns out, torch.compile does try and pad your matmuls to have the right shape! See the code here, or try this benchmark.

import torch
torch.set_default_device('cuda')
from triton.testing import do_bench

def f(a, b):
    return torch.mm(a, b)

a = torch.randn(4096, 4096, dtype=torch.bfloat16)
b = torch.randn(4096, 4095, dtype=torch.bfloat16)
print("eager: ", do_bench(lambda: f(a, b)))
cf = torch.compile(f)
print("compiled: ", do_bench(lambda: cf(a, b)))
>> eager: 1.4077268838882446
>> compiled: 0.6021425127983093

However, there are still limitations that mean it makes sense for users to manually pad their shapes.

For one, padding requires a full copy! Although torch.compile can often fuse this into a preceding op, in the case where the matrix being padded comes from the input (like a weight matrix), there’s no way to avoid this copy.

Second, resolving wave quantization is far more difficult.

Conclusion

Overall, I hope the topic of "how do I squeeze the most out of my matmuls" is an interesting one. There's still many more intricacies in matmul perf that I didn’t have the time to get to, as well (I’m sure) many more intricacies that I don’t know! Here’s the main code to replicate the results.

Also, here’s some quiz questions to test your understanding! I will publish a brief explanation of the answers at some later point.

Quiz Questions

1: Let's say I have a [M x K] @ [K x N] matmul. Which one of these configurations will have the best perf? Think about the actual ramifications of tiling! Both matrices are in row-major layout (i.e. K and N are the innermost dimensions)
A: M=2047, K=N=2048
B: K=2047, M=N=2048
C: N=2047, M=K=2048

2: Let’s say I have an A100 with 108 SMs, and I want to benchmark a number of matmuls with no wave quantization. How would I go about constructing the shapes for these matmuls?

3: Based off this post, would you expect that making your batch size a power of 2 leads to more efficient performance?

4: Similar to Question 1, let’s say we have a A: [M x K] @ B: [K x N] matmul. However, now, A is in column-major (i.e. torch.randn(K, M).t()) while B is still row-major. What is the best configuration now?

A: M=2047, K=N=2048
B: K=2047, M=N=2048
C: N=2047, M=K=2048

5: Let’s say that we have this code.

A = torch.randn(4096, 4096)
B = torch.randn(4096, 4096)
B = B[:, :4095] # B now has shape [4096, 4095]

Would you expect that we have good performance on a matmul between A and B?

Solutions can be found below

Featuring:

• (!) no custom kernels

• int8 and tensor-parallelism support

• still very simple (<150 LOC to support)

• faster decoding than any (non-Groq) API endpoint, at up to 220 tok/s/user with A100s.

Bonus: Running on a H100 node, we can get closer to 300 tok/s/user!

We thought it might be interesting to talk a bit about why this is a bit tricky and how we solved it.

Mixture of Experts (MoE) vs. Dense Transformers

Unlike Llama, Mixtral is a sparse architecture. The main idea here is that instead of a single dense layer with eight thousand parameters, we split it into 8 “experts”, where each expert is a dense layer with only one thousand parameters. Then, depending on some dynamic information, we only “activate” 2 out of the 8 experts for each layer.

Morally, this is easy to implement for each token.

Crucially, note that out of the 8 experts that constitute our weights, 6 of them do nothing for any token, making mixture of experts a “sparse” model.

For those who have experience with performant PyTorch, you may be wincing. Using a tensor to index into a python list is a cardinal performance sin (it induces a cuda sync, where the CPU waits for the GPU).

This is the difficulty with running Mixtral efficiently.

The runtime advantage of MoE comes from dynamic sparsity. But if this dynamism isn't handled efficiently, you might end up slower than if you had no sparsity to begin with.

Moving the Dynamism “onto” the GPU

Luckily, there’s another option that works well for BS=1 decoding. Instead of doing the indexing in Python, let’s do the indexing on the GPU. In other words, let’s do the indexing using a “gather” operation.

A gather operation occurs when you decide to load from a tensor using another tensor. In PyTorch, this is often done using what’s called “advanced indexing”. For example:

primes = torch.tensor([2, 3, 5, 7, 11])
b = torch.tensor([1, 3, 0])
primes[b] # tensor([3, 7, 2])

So, putting it all together, our full MoE layer looks morally like this.

The full FFN layer looks a bit more different since we need to handle multiple tokens as well as multiple FFNs in a row, but it’s morally the same idea. Note: this is the primary implementation difference between regular dense transformers and Mixture of Experts!

This implementation approach has two main advantages - it doesn't require any synchronizations and only uses weights required for computation. However, if we run this normally, we have another issue. Both the gather operation and the actual linear layer itself requires us to touch DRAM with all of the weights we’re using. This is a factor of 3 slowdown.

Luckily, PyTorch now has a compiler.

Torch.compile to the rescue!

Torch.compile can fuse the gather + gemv into one kernel, allowing us to obtain our theoretical speedups.

If you’re curious to look at the Triton kernel generated by torch.compile, you can see it here.

Concretely, this is the indirect access/gather. Bolding added to emphasize the main operations involved.

tmp0 = tl.load(in_ptr0 + (r2 + (4096*x1)), None, eviction_policy='evict_last').to(tl.float32)
tmp1 = tmp0.to(tl.float32)
tmp3 = tmp2 + 8
tmp4 = tmp2 < 0
tmp5 = tl.where(tmp4, tmp3, tmp2)
tmp6 = tl.load(in_ptr2 + (r2 + (4096*(x0 % 14336)) + (58720256*tmp5)), None, eviction_policy='evict_first')
tmp7 = tmp6.to(tl.float32)
tmp8 = tl.load(in_ptr3 + ((14336*tmp5) + (x0 % 14336)), None, eviction_policy='evict_first').to(tl.float32)

Let’s also validate the performance with a benchmark.

def cuda_indexing(W, score_idxs, x):
    return W[score_idxs] @ x

def python_indexing(W, score_idxs, x):
    return W[score_idxs[0]] @ x, W[score_idxs[1]] @ x

W = torch.randn(E, D, D)
x = torch.randn(D)
score_idxs = torch.tensor([3, 5])

End to End Benchmarks

Combining this altogether into an E2E benchmark, we see that for int8 on a single A100, we run at 98 tok/s. Note that if this were a dense model, we would effectively be running at 4.55 TB/s of bandwidth, which is higher than the theoretical limit!

Of course, combining it with tensor-parallelism, we can get up to 280 tok/s! Note that this is all for BS=1, so this is “tok/s/user for only output tokens”.

Considering that this is faster tokens/s than any (non-Groq) API provider, we think this is pretty impressive.

Moreover, since we codegen into Triton, we should also be able to run on AMD as well. We’ll update this post when we get those results.

Of course, we will also mention the typical gpt-fast caveats. This is optimized for latency and not throughput. In this particular case, the strategy we use for BS=1 codegen scales very poorly to larger batch sizes.

Nevertheless, we think that this continues to demonstrate the gpt-fast ethos. Simple, native PyTorch, and very fast!