How To Write C++ Extensions in 2026 - Jane Xu, Meta & Mikayla Gawarecki, Meta

This talk discusses how to create a C++ custom extension for PyTorch in 2026, emphasizing the process and benefits of writing such extensions. The speakers, Jane and Michaela from Meta, outline the reasons to extend PyTorch, such as the need for specific operations not available in the standard library. They demonstrate writing a fused multiply add operation while utilizing the stable ABI from PyTorch to ensure compatibility across versions. The session covers essential APIs, memory management techniques, and registration with the PyTorch dispatcher. They also touch on ABI stability, compiler flags, and methods to compile and use extensions efficiently within the PyTorch ecosystem.

All right, good morning. Hi, everyone. I'm Jane. I work on PyTorch at Meta. >> And I'm Michaela, and I also work on >> And today we're going to discuss how to write a good C++ custom extension or custom op in 2026. Um if you have any questions, just like raise your hand. It's a small room, so like don't worry about it. We also have like 3 extra

minutes of buffer time, so don't feel like, you know, you're interrupting anything. But anyway, how do we write a good C++ extension in 2026? Well, first we should ask, why should we write a C++ extension at all? And that's two questions. The first question is, when would you want to extend PyTorch? And the second question is, why C++? So, you may want to extend PyTorch if you

are trying to use PyTorch, and then you're like, wait, I have a cool op that I want to write myself that's not supported in our 2008 and operators. This happens more often than you think. And um when you want to do that, you usually also want to still use the rest of the torch system. For example, you still want autograd, you still want PT2, you want to

be able to torch compile and work with the rest of the PyTorch ops. And you could technically extend PyTorch with a Python extension. It's like very easy to do that as well with like cute DSL or Triton or whatever you want or even in other torch ops. But you might want to write C++ for several reasons. Um they're not all listed here, but for example, you might

want to target specific hardware or you know, use certain heuristics that would be better for perf for your particular use case. And or maybe you just really like writing C++. Uh that is very valid. You can you can do that, too. Um So now that we've settled on, hey, we're going to write a C++ extension, let's actually walk through writing one together. Today we're going to use

a fused multiply add. This is This is just an example. This is not like a real-life use case. Um but semantically in Python, suppose that you're trying to add two you're trying to multiply two tensors and then add a float to it. So, here you have tensor A and B, and then you have float C. You're adding C to A * B. For simplicity, we are going

to assume that no matter what A and B are, we are going to return an out-of-place contiguous result. So, even if A is not contiguous, we're going to allocate new memory to return a new contiguous answer. And yeah, uh this is for illustration purposes. We know like views multiply add is not a completely new feature, uh but follow along. Okay, so before we start writing, you might

wonder, "Okay, how do you even use like the the APIs? What what APIs are there even?" Um, and this diagram is from my ABI stable talk from last PTC, where I'm not going to go into all the details, but if you're curious about how this was formulated, there is a QR code for later time. Um, but you can imagine that libtorch is this like bottom mountainous thing

at here that's like jagged. And what the jagged means is that libtorch is still under development. We want to keep improving it, so like it can move around. But, that means that if you depend on libtorch, you have to move along with the with the jaggedness. So, if libtorch changes, you have to change. You have to update, rebuild your extension. However, as of last year and now,

we have a stable API. Well, we have a C shim stable ABI from for it from which we built high-level C++ wrappers So, you can imagine the C shim handles all the jaggedness for you. So, we're promising that even if libtorch is changing under the hood, we have a C shim that will stay the same. So, if you use any of the APIs that depend on the

C shim, you're good. You don't need to update when libtorch updates. Um, so to just walk into it, there are like I would think of this as three chunks. You have the C shim, which is in C, and as you can imagine if you've ever worked with C versus C++, it's not as fun. But, because we know it's not as fun, we added some high-level C++ wrappers.

So, for example, in libtorch, you are familiar probably with ATen tensor. Well, in the torch C source stable, we have torch stable tensor, which looks very semantically similar to tensor. So, your user experience is supposed to be still okay after this. And then on the left, you might be like, what is that column of green that like doesn't depend on libtorch? It is what we refer to

as our header-only APIs. So, some examples are like dtype or device type, things that don't actually depend on the libtorch.so. And these are header-only because they're fully implemented with headers, therefore they're standalone, and at the same they also do not move when libtorch moves. Okay. So, whenever you can, you should use our ABI st- uh stable libtorch API. Oh, that is a mouthful. Um and why should

you do that? Because just like I mentioned earlier, if libtorch changes, you don't want to have to re-update your extension every time. That's kind of a lot of wheels. What it means is for PyTorch 2.10 to 2.11 to 2.12, you only have to build one wheel versus three wheels. Um note uh this is for 2.10 plus, which we're on 2.11 and we're about to like start 2.12,

so you could start using this. And note that it is particularly a limited subset. So, it's not meant to replace libtorch, it's not meant to like do more than uh what we suggest, which is where using it for custom ops. Um you can always fall back to the non-stable version, but you might miss out on some features. And so, if you feel FOMO and you're like, why

is this thing not ABI stable yet? Just like open an issue, we will get to it, we'll respond to you. It is still under development, so don't be shy. Okay, well, yeah, why ABI stability? Um, this is to kind of reiterate the earlier point of Normally, when you are building your extension, you have to think about like the Python versions and the torch versions. And you can

imagine this box going in other dimensions. Like, you might have to build for CUDA versions or like different architectures. So, it's bigger than just two axes. But, we're going to think about these two today. Um, because with libtorch stability, you can build Well, here you have to build like one binary per version. And you can turn that into for all the versions. So, that delineates one of

your axes. So, instead of something like O of M * N, you just get O of N. We'll get into more of this later, too. All right. So, if you're like, "How do I make sure that I'm using the right headers?" Well, we made it easy for you. There are only like three things you can include. Okay, four things, but that's cuz we have two shims. But,

mostly three things. And they they correspond to the categories I was saying earlier. So, there's the C shims, which is at the bottom. There's your header only, like D types, uh, and device type. And then there's stuff like stable tensor, stable device that live in the torch C source stable. So, normally you will see code if if you see code that looks like the left where you

in torch include things like torch library.h, you can't do that anymore. If you do that, it will not be good. You you won't be you won't be But, we have alternatives on the right-hand side that you can use that should suffice for most of your use cases. And we do enforce this at build time with this torch target version flag. Michaela's going to talk more about that

later, so sit tight. But, note that if you pass this flag, we will yell at you at compile time if you do something wrong, too. So, there are many We we try to make it so that you get signal as soon as possible. Okay, so let's actually write the code. Uh remember our reference in Python looks like that. We're going to start with the schema. Uh with

most ops you have to define a schema, so it's good to think about that. Here it's chill. We're taking in two tensors. Note that and returning a tensor. So, note that I'm using the torch stable tensor right off the bat. And then we're using double for float. That is just how ATen handles it. Don't don't ask me too many questions there. Um and then we're going to

return a Uh yes. Also, there are docs for all of the things. I do want to call that out. Uh I don't have a QR code because you should actually just be able to Google this. And if you're Googling it and you can't find it, please let us know because that's bad. Uh Anyway. Okay. So, the first easy thing we can do is just add some checks.

And here I mostly want to point out that you that this should look similar to you and familiar. Like it's not supposed to look like alien code. Like torch check is just standard torch check in the stable world. You can still call things like dot sizes, dot scalar type. You can do You can get like floats and CUDA. So, here we're just checking that A and B

are the same size, that they are both floats, and they're both on CUDA. Um yeah. And if you're like, "How do I know what APIs are available?" They're there again in our C++ docs. Um And Oh, yeah. That's Let's keep going. Okay. So, we're going to write some CUDA. This is very basic CUDA. This is not what like the CUDA looks like in PyTorch, but this is

pretty pretty good. Okay. So, we because we are assuming that all the tensors will become contiguous at this point, the math is pretty chill. You have your tensor that's going to look like a float star. Um and then you have your float that's going to look like float. And then we're going to write into a float star result. And literally the math is just like A at

a certain index times B at a certain index plus C. This part should not be too surprising. But of course with CUDA, oh wait, yeah. Uh there's some space here and you don't the for the purpose of of this presentation you don't have to understand this, but uh this is why we have in numel, which is the number of elements for indexing into blocks, but it's okay.

Wait, let me just let me just go next so it makes more sense. Um when we call the CUDA kernel, we are calling it with the CUDA launch params of like grid size, block size, shared memory, and the stream. And so based on the grid size and block size, that is what is going to get like that that that's what this whole block index block dim thread

index stuff is. But you can just think of it as doing math on a contiguous array of three three uh tensors. Um okay, what do you notice here? So notice that while we have arguments, we're assuming block size is 256 and we have all those CUDA things set up, we don't actually have what like numel, A pointer, B pointer, C. Um well actually no, we have C,

but we don't have a result pointer yet. So let's actually set up those inputs. So I've recorded all the things we need in that comment. We're just going to start getting into it. So first remember that we have to make everything contiguous. The main thing to note here is that the stable ABI has ops. So like if you wanted to call A10 ops, you should be able

to with the stable operators we provide. And we're going to call contiguous on them. You can call things like empty like. We see that everywhere in most custom ops, you probably want something like that. If you want more ops, just let us know, uh but they should be supported. And note that we're going to get the pointer of A and B with const data pointer since we're

reading into it and then we're going to call it mutable data pointer on the result since we're writing into it. Um the key thing to notice here is that hey, all these APIs that you probably want are supported. And if they're not, you can let us know. And lastly, we still need stream and Numel. Um, Numel is just Numel, that's chill. The stream API here looks a

little messy. I agree. This is kind of to exemplify if you wanted it to use our shim. It does not look great cuz it's in C. So, it's like you know, not very pretty. We are also going to make this better. Like this is actually the only way you can do it today because uh lolz. But uh we are going to make this Well, Michaela said she

would make this better, so you can keep her accountable. Um, and and the this is just how you can get a stream today. If you're like, this is so much code, what what's going on? Like we're done here. Um, but all this code is also just in our example repo at the bottom there uh in extension.cpp. So, if you're like, this is kind of a lot to

follow along, you don't have to remember everything. This is already in our repo somewhere. Um, but this is all the code we just wrote. And now we're done with the C++ maybe. >> Okay, cool. So, let's recap what Jean talked to us about. First of all, she introduced us to the three classes of PyTorch stable ABI headers. And then she walked us through an example of how

to write a simple my model at CUDA kernel using the APIs from these headers. So, what haven't we talked about yet? We haven't talked about how to register this op to the dispatcher such that it can interact with the PyTorch dispatcher in an ABI stable manner. And we haven't talked about how to compile your extension yet. So, let's first talk about registering it to the PyTorch dispatcher.

Essentially, all you have to do is to register the schema and the operator implementation using the stable torch library macro. Um, for those people who are familiar with the existing ABI unstable way of registering extensions to torch, you would probably notice two things. Um, first of all, the ABI unstable header is called torch library, but this one is called stable torch library. And second the implementation that

is passed to stable torch library impel has to be boxed with the torch box macro. After you do these two things, you and you compile your extension, you can then call it from Python as such. But before we go into the Python side of things, let's talk a little bit more about torch box. Okay, so why do we need torch box? What does it do? Recall earlier

our function took in tensor tensor float. Torch box always converts the function into the following box function format, which takes these three arguments. The first argument is a stable I value star stack, which is expected to take in like a stack of stable I values from the dispatcher. The second argument is the number of arguments as parsed from the function schema that's registered. And the third argument

is the number of outputs, also as parsed from the function schema registered. The first thing it does is that it's going to use the two function, which converts from stable I value to the appropriate object type from the schema. And then after that, it's going to apply your function that you have registered or your kernel that you have registered. And then after that, it will take the

output from your kernel and then use the from function, which converts back from stable sorry, back from the object type to stable I value. And then place this object back onto the stack. Okay, so I mentioned stable I value multiple times on the previous slide, but I didn't explain what a stable I value actually is. So let's talk a little bit about that. Essentially, a stable I

value is the ABI stable representation of the object. So notice that the inputs to your kernel as well as the outputs to your kernel are actually crossing the libtorch ABI boundary. you could imagine that your extension might have been compiled in let's say like torch 2.10, and there could be some object A which has like a certain representation in torch 2.10. And then, at run time, you're

actually running your extension with let's say torch 2.11. So, um, your dispatcher is passing in the arguments with like the libtorch 2.11 version, which might have another version of that object A. Uh, so, stable ABI value is like sort of this like common representation, and two and from are these like translation layers that translate to and from this ABI stable representation. Um, yeah. So, now that we

talked about torch box, we can go back to like, how do you compile your extension? Um, so, this slide shows like an example of how you use torch util CPP extension, as well as like setuptools to compile your extension. Um, the two things that I want to highlight here are the two compiler flags that we have passed, uh, pile limited API and torch target version. Jane mentioned

torch target version earlier. Uh, so, let's get a little bit into So, what is torch target version? Um, the format that torch target version expects is shown on this slide. Essentially, you're supposed to pass in a certain major and minor torch torch version. Uh, so, for example, if you wanted to target 2.10, this is what the compiler flag would look like. Um, and what purposes does it

serve? Jane touched on this a little bit earlier. First of all, it will ban the extension from including any libtorch unstable headers. So, like in this example, if you accidentally like use ATen CUDA context and some API from there, you're going to get a loud and angry compilation error that tells you, "Don't do that." Okay, second of all, um, kind of as the name suggests, it allows

you to select a minimum torch version that the compiled extension runs with. Um, note that this does not have to be the same as the build time version of torch. So, for example, you could download a 2.12 torch nightly, you could build the extension using the 2.10 torch target version flag, and then your compiled extension is guaranteed to run with any torch nightly or stable release that's

greater than or equal to torch 2.10. Um also note that this feature doesn't like backdate that much. So, like don't expect to be able to pass like torch 2.4 as your target version and for your compiled binary to run on torch 2.4. Okay. Um let me talk a little bit about how torch target version works. Um I'll talk about this from the perspective of the three classes

of libtorch uh ABI stable headers that Jane mentioned So, to recap, we have the C shim. Um we have the ergonomic wrappers in torch C source stable, and those are usually like doing things like error handling or memory management, which is like a little bit more nice in the C And last of all, we have the libtorch free torch header only headers. So, on the left of

this slide, we have like an example API from the C shim, which is like torch from blob. Um let's say that this was introduced in version 2.11. There will always be if dev macros around this um shim that prevent it from being exposed unless you target a version that's greater than or equal to the version that it was exposed in. And then on the right of the

slide, we have the nice um wrapper API in torch C source stable. In this example, it's very simple. It just like is a from blob API that calls the C shim, torch from blob. So, um the APIs in torch C source stable will always be version based on the max version of any shim it uses. So, in this case, it's just like the same 2.11 gates. And

conspicuously missing from this slide is torch header only. So, torch header only has no dependency on libtorch. So, um the version of an object that's available at build time is always the version that is compiled into Um there's one caveat here, which is that if that object crosses the ABI boundary, there would actually be some implication on ABI stability. But then that will be like appropriately Sorry,

first of all, that is likely not to happen because we sort of guarantee to freeze torch header only. But second of all, if it does happen, the firm into um conversions will like handle that translation properly. Okay, so why ABI stability? Going back to this diagram Jean showed, you can build a single binary that targets multiple Sorry, that works with multiple torch versions above a certain version.

But there's another flag that I mentioned just now, which is Py_LIMITED_API. So, let's talk a little bit about that. There's actually a second dimension of ABI stability that we can also like think about for our extensions, which is CPython ABI stability. Um this Py_LIMITED_API flag is provided by CPython. Um and you can pass it to help verify your extension only uses the CPython stable limited API. Um

note that for example, headers like PyBind are not CPython ABI stable. So, like if you use them, your extension might not be able to be CPython agnostic. Um yeah, and one like weirdish thing is that you need to pass this flag in three different places for your extension to actually compile properly with like Python agnosticism. So, yeah, this flag could be useful for future reference if you're

trying to do this. Okay. So, back to our diagram. Um now we have one binary. We went from m * n binaries, and now we have one single binary that works across multiple torch and Python versions. Okay, now let me go back to calling your op in Python. we recommend this way of like finding your SO file and then using torch.ops.load_library to load the SO file. Uh

the reason why we would recommend this is because it doesn't use any CPython APIs. So, um that's good for CPython ABI stability. Um and after you do this, your op is now available from Python and you can call it as torch.ops.extension_cpp_stable.mymullet. Um yeah, there are other ways to expose C++ ops, but we recommend this way to achieve CPython ABI stability, as I already said. Yeah. Um okay,

now I've talked about the two dimensions of ABI stability, and then let me just talk a little bit about some other nice utilities that torch.library provides that you can use with your like kernel that you've written. So, first of all, there's torch.library.register_fake. Um we recommend that you use this from Python to register your faker meta kernel for like your op to compose with torch.compile. And second of

all, another like very useful thing is torch.library.opcheck. It will check certain some things, like for example, whether the schema of your function that you've registered actually matches the schema of the kernel, or some other things, like whether your fake tensor implementation is actually correct. Um note that opcheck testing is not exhaustive, so it's still good to write your own tests. yeah. So, actually that brings us to

the conclusion of our talk. Um as Jane mentioned earlier, if you want to actually look at that mymullet example, it's on this uh extension_cpp_stable PyTorch repository, and I think that is the QR code for it, so like you can scan that if you're interested. And there are also several libraries that are already on board the torch stable ABI. So, we have flash attention three, xformers CUDA, torch

audio, torch io CUDA more recently, and VLM CUDA is also in progress. Cool. So, yeah, that brings us to the end of our talk. And if you have any questions, feel free. >> Does anyone have any questions? Wow, okay. Very very very exhaustive talk, I suppose. No, just kidding. Um I would say that that link is the tutorial, so it's even better than the repo. It also

does link to the repo. And then also, all of those legit use cases in the wild are like maybe even better examples for your real-life use case. So, like, I would actually go to those if you're like, "Okay, this this MoLA thing is too easy." If you wanted something more legit, go go go there. Um and these slides, I think, are supposed to be online. So, if

you didn't pay attention, you can you can review there. All right, thanks, guys.