deepseek-ai / deepseek-ai/FlashMLA
Make FlashMLA a libtorch and cpython stable extension
- Dominant language
- C++
- Stars
- 12.9k
- Forks
- 1.2k
- Avg merge
- 4h 20m
- Merged PRs (30d)
- 2
Description
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:
1. 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.
2. 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)
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by reviewing the five named kernels—sparse_decode_fwd, dense_decode_fwd, sparse_prefill_fwd, dense_prefill_fwd, and dense_prefill_bwd—and the linked PyTorch stable ABI documentation. Confirm the maintainers want this migration, then assess dispatcher registration and stable ABI equivalents; done means the kernels build and release as a stable extension across the targeted PyTorch and CPython versions.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, pytorch
- Domain
- machine-learning
- Issue type
- Refactor
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 30/100