deepseek-ai / deepseek-ai/FlashMLA

Make FlashMLA a libtorch and cpython stable extension

Open
#180 0 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.