Hello! I work on PyTorch and in particular I've been focusing on solving the ecosystem problem where PyTorch custom extensions must rebuild for new PyTorch versions (as libtorch.so guarantees no ABI stability). Hence, we've been developing a limited stable ABI subset of libtorch.so to cover the most common custom op use cases. See https://docs.pytorch.org/docs/stable/notes/libtorch_stable_abi.html for more deets.
Motivation
It may be of community interest to have FlashMLA be detangled from libtorch by migrating its kernels to use ABI stable APIs, especially as it's a well-used building block by several custom kernel libs like vLLM. This way, you'd only have to build and release one FlashMLA wheel across multiple PyTorch versions.
Proposal
I've done a glance through the 5 kernels (sparse_decode_fwd, dense_decode_fwd, sparse_prefill_fwd, dense_prefill_fwd, and dense_prefill_bwd) and what libtorch APIs they use and I believe we have stable ABI equivalents for them from PyTorch 2.10+. My proposed plan would involve:
- Moving from using PYBIND to registering with the torch dispatcher through TORCH_LIBRARY --> is there a reason we shouldn't do this? this would also unlock CPython ABI stability, meaning you'd be able to build 1 wheel across multiple CPython versions.
- Migrating the kernels to use the stable ABI (at::Tensor -> torch::stable::Tensor, TORCH_CHECK -> STD_TORCH_CHECK).
Questions for the maintainers
mainly: Would being libtorch stable be of interest to the library?
How come FlashMLA doesn't register through torch.library? Some benefits would include detangling from CPython, unlocking torch.compile support + other custom op benefits (https://docs.pytorch.org/tutorials/advanced/cpp_custom_ops.html)
Hello! I work on PyTorch and in particular I've been focusing on solving the ecosystem problem where PyTorch custom extensions must rebuild for new PyTorch versions (as libtorch.so guarantees no ABI stability). Hence, we've been developing a limited stable ABI subset of libtorch.so to cover the most common custom op use cases. See https://docs.pytorch.org/docs/stable/notes/libtorch_stable_abi.html for more deets.
Motivation
It may be of community interest to have FlashMLA be detangled from libtorch by migrating its kernels to use ABI stable APIs, especially as it's a well-used building block by several custom kernel libs like vLLM. This way, you'd only have to build and release one FlashMLA wheel across multiple PyTorch versions.
Proposal
I've done a glance through the 5 kernels (
sparse_decode_fwd,dense_decode_fwd,sparse_prefill_fwd,dense_prefill_fwd, anddense_prefill_bwd) and what libtorch APIs they use and I believe we have stable ABI equivalents for them from PyTorch 2.10+. My proposed plan would involve:Questions for the maintainers
mainly: Would being libtorch stable be of interest to the library?
How come FlashMLA doesn't register through torch.library? Some benefits would include detangling from CPython, unlocking torch.compile support + other custom op benefits (https://docs.pytorch.org/tutorials/advanced/cpp_custom_ops.html)