pybind->TORCH_LIBRARY; get ABI stable with CPython by janeyx99 · Pull Request #193 · deepseek-ai/FlashMLA

janeyx99 · 2026-07-02T15:33:47Z

Fixes #180

The first commit migrates the library from pybind -> TORCH_LIBRARY. This allows the library to be CPython agnostic, meaning you can build one wheel for multiple python versions (3.10+ here). Then, the second commit does the mechanical work of migrating APIs to the stable APIs.

After the first commit, there are 70 unstable APIs used by the .so:

(gb-flashmla-pt211) ➜  FlashMLA git:(abi-stable-migration) ✗   torch-abi-audit flash_mla/cuda.abi3.so  
Package: cuda.abi3.so
  Root: /home/janeyx/repos/FlashMLA/flash_mla
  Torch ABI:   UNSTABLE
  CPython ABI: n/a
  Extensions:  0
  Bundled libs: 1
  -- bundled libs --
    [UNSTABLE] [abi3-tagged-no-capi   ] cuda.abi3.so  (stable_shim=0, unstable=70)

After the second, there are none!!

(gb-flashmla-pt211) ➜  FlashMLA git:(abi-stable-migration) ✗ torch-abi-audit flash_mla/cuda.abi3.so  
Package: cuda.abi3.so
  Root: /home/janeyx/repos/FlashMLA/flash_mla
  Torch ABI:   STABLE
  CPython ABI: n/a
  Extensions:  0
  Bundled libs: 1
  -- bundled libs --
    [STABLE  ] [abi3-tagged-no-capi   ] cuda.abi3.so  (stable_shim=63, unstable=0)

Test Plan

After commit 1:

python tests/test_flash_mla_sparse_decoding.py
...
4748/4748 correctness cases passed
TFlops     geomean: 329.9

After commit 2 (showing that perf diff is insignificant):

python tests/test_flash_mla_sparse_decoding.py
...
4748/4748 correctness cases passed
TFlops     geomean: 329.8

After both commits (ran twice):

python tests/test_flash_mla_sparse_prefill.py
...
All 617 cases passed!

test_fmha_sm100.py also had no failures:

Details

# commit 1:
python tests/test_fmha_sm100.py     
b=2, mean_sq=4096, mean_sk=4096, varlen=False, h=128, h_k=128, d=128, dv=128, causal=False, has_bwd=True, check_correctness=True
1.607 ms, 1369 TFLOP/s, name: fwd
4.978 ms, 1104 TFLOP/s, name: bwd
b=2, mean_sq=4096, mean_sk=4096, varlen=False, h=128, h_k=128, d=128, dv=128, causal=True, has_bwd=True, check_correctness=True
0.966 ms, 1139 TFLOP/s, name: fwd
3.017 ms, 911 TFLOP/s, name: bwd
b=2, mean_sq=4096, mean_sk=4096, varlen=False, h=128, h_k=128, d=192, dv=128, causal=False, has_bwd=True, check_correctness=True
1.765 ms, 1557 TFLOP/s, name: fwd
9.125 ms, 783 TFLOP/s, name: bwd
b=2, mean_sq=4096, mean_sk=4096, varlen=False, h=128, h_k=128, d=192, dv=128, causal=True, has_bwd=True, check_correctness=True
1.106 ms, 1243 TFLOP/s, name: fwd
4.736 ms, 755 TFLOP/s, name: bwd
b=2, mean_sq=4096, mean_sk=4096, varlen=False, h=32, h_k=4, d=128, dv=128, causal=False, has_bwd=False, check_correctness=True
0.431 ms, 1276 TFLOP/s, name: fwd
b=2, mean_sq=4096, mean_sk=4096, varlen=False, h=32, h_k=4, d=128, dv=128, causal=True, has_bwd=False, check_correctness=True
0.253 ms, 1085 TFLOP/s, name: fwd
b=2, mean_sq=4096, mean_sk=4096, varlen=False, h=32, h_k=4, d=192, dv=128, causal=False, has_bwd=False, check_correctness=True
0.448 ms, 1533 TFLOP/s, name: fwd
b=2, mean_sq=4096, mean_sk=4096, varlen=False, h=32, h_k=4, d=192, dv=128, causal=True, has_bwd=False, check_correctness=True
0.283 ms, 1213 TFLOP/s, name: fwd
b=2, mean_sq=4096, mean_sk=4096, varlen=True, h=128, h_k=128, d=128, dv=128, causal=False, has_bwd=True, check_correctness=True
2.120 ms, 1171 TFLOP/s, name: fwd
6.436 ms, 964 TFLOP/s, name: bwd
b=2, mean_sq=4096, mean_sk=4096, varlen=True, h=128, h_k=128, d=128, dv=128, causal=True, has_bwd=True, check_correctness=True
1.323 ms, 1105 TFLOP/s, name: fwd
4.133 ms, 884 TFLOP/s, name: bwd
b=2, mean_sq=4096, mean_sk=4096, varlen=True, h=128, h_k=128, d=192, dv=128, causal=False, has_bwd=True, check_correctness=True
2.343 ms, 1324 TFLOP/s, name: fwd
9.597 ms, 841 TFLOP/s, name: bwd
b=2, mean_sq=4096, mean_sk=4096, varlen=True, h=128, h_k=128, d=192, dv=128, causal=True, has_bwd=True, check_correctness=True
1.485 ms, 1231 TFLOP/s, name: fwd
6.169 ms, 770 TFLOP/s, name: bwd
b=2, mean_sq=4096, mean_sk=4096, varlen=True, h=32, h_k=4, d=128, dv=128, causal=False, has_bwd=False, check_correctness=True
0.514 ms, 1206 TFLOP/s, name: fwd
b=2, mean_sq=4096, mean_sk=4096, varlen=True, h=32, h_k=4, d=128, dv=128, causal=True, has_bwd=False, check_correctness=True
0.345 ms, 1059 TFLOP/s, name: fwd
b=2, mean_sq=4096, mean_sk=4096, varlen=True, h=32, h_k=4, d=192, dv=128, causal=False, has_bwd=False, check_correctness=True
0.555 ms, 1397 TFLOP/s, name: fwd
b=2, mean_sq=4096, mean_sk=4096, varlen=True, h=32, h_k=4, d=192, dv=128, causal=True, has_bwd=False, check_correctness=True
0.378 ms, 1208 TFLOP/s, name: fwd
b=2, mean_sq=8192, mean_sk=8192, varlen=False, h=128, h_k=128, d=128, dv=128, causal=False, has_bwd=True, check_correctness=False
7.031 ms, 1251 TFLOP/s, name: fwd
21.337 ms, 1031 TFLOP/s, name: bwd
b=2, mean_sq=8192, mean_sk=8192, varlen=False, h=128, h_k=128, d=128, dv=128, causal=True, has_bwd=True, check_correctness=False
3.762 ms, 1169 TFLOP/s, name: fwd
13.013 ms, 845 TFLOP/s, name: bwd
b=2, mean_sq=8192, mean_sk=8192, varlen=False, h=128, h_k=128, d=192, dv=128, causal=False, has_bwd=True, check_correctness=False
8.466 ms, 1299 TFLOP/s, name: fwd
30.886 ms, 926 TFLOP/s, name: bwd
b=2, mean_sq=8192, mean_sk=8192, varlen=False, h=128, h_k=128, d=192, dv=128, causal=True, has_bwd=True, check_correctness=False
4.089 ms, 1345 TFLOP/s, name: fwd
16.598 ms, 861 TFLOP/s, name: bwd
b=2, mean_sq=8192, mean_sk=8192, varlen=False, h=32, h_k=4, d=128, dv=128, causal=False, has_bwd=False, check_correctness=True
1.543 ms, 1425 TFLOP/s, name: fwd
b=2, mean_sq=8192, mean_sk=8192, varlen=False, h=32, h_k=4, d=128, dv=128, causal=True, has_bwd=False, check_correctness=True
0.857 ms, 1283 TFLOP/s, name: fwd
b=2, mean_sq=8192, mean_sk=8192, varlen=False, h=32, h_k=4, d=192, dv=128, causal=False, has_bwd=False, check_correctness=True
1.744 ms, 1576 TFLOP/s, name: fwd
b=2, mean_sq=8192, mean_sk=8192, varlen=False, h=32, h_k=4, d=192, dv=128, causal=True, has_bwd=False, check_correctness=True
0.956 ms, 1437 TFLOP/s, name: fwd
b=2, mean_sq=8192, mean_sk=8192, varlen=True, h=128, h_k=128, d=128, dv=128, causal=False, has_bwd=True, check_correctness=False
8.136 ms, 1221 TFLOP/s, name: fwd
24.892 ms, 997 TFLOP/s, name: bwd
b=2, mean_sq=8192, mean_sk=8192, varlen=True, h=128, h_k=128, d=128, dv=128, causal=True, has_bwd=True, check_correctness=False
5.592 ms, 1046 TFLOP/s, name: fwd
15.553 ms, 940 TFLOP/s, name: bwd
b=2, mean_sq=8192, mean_sk=8192, varlen=True, h=128, h_k=128, d=192, dv=128, causal=False, has_bwd=True, check_correctness=False
10.012 ms, 1240 TFLOP/s, name: fwd
36.787 ms, 877 TFLOP/s, name: bwd
b=2, mean_sq=8192, mean_sk=8192, varlen=True, h=128, h_k=128, d=192, dv=128, causal=True, has_bwd=True, check_correctness=False
5.700 ms, 1282 TFLOP/s, name: fwd
22.357 ms, 850 TFLOP/s, name: bwd
b=2, mean_sq=8192, mean_sk=8192, varlen=True, h=32, h_k=4, d=128, dv=128, causal=False, has_bwd=False, check_correctness=True
1.855 ms, 1338 TFLOP/s, name: fwd
b=2, mean_sq=8192, mean_sk=8192, varlen=True, h=32, h_k=4, d=128, dv=128, causal=True, has_bwd=False, check_correctness=True
1.151 ms, 1270 TFLOP/s, name: fwd
b=2, mean_sq=8192, mean_sk=8192, varlen=True, h=32, h_k=4, d=192, dv=128, causal=False, has_bwd=False, check_correctness=True
2.064 ms, 1503 TFLOP/s, name: fwd
b=2, mean_sq=8192, mean_sk=8192, varlen=True, h=32, h_k=4, d=192, dv=128, causal=True, has_bwd=False, check_correctness=True
1.275 ms, 1433 TFLOP/s, name: fwd

# Commit 2:
(gb-flashmla-pt211) ➜  FlashMLA git:(abi-stable-migration) ✗ python tests/test_fmha_sm100.py           
b=2, mean_sq=4096, mean_sk=4096, varlen=False, h=128, h_k=128, d=128, dv=128, causal=False, has_bwd=True, check_correctness=True
1.569 ms, 1402 TFLOP/s, name: fwd
5.419 ms, 1014 TFLOP/s, name: bwd
b=2, mean_sq=4096, mean_sk=4096, varlen=False, h=128, h_k=128, d=128, dv=128, causal=True, has_bwd=True, check_correctness=True
0.968 ms, 1136 TFLOP/s, name: fwd
3.023 ms, 909 TFLOP/s, name: bwd
b=2, mean_sq=4096, mean_sk=4096, varlen=False, h=128, h_k=128, d=192, dv=128, causal=False, has_bwd=True, check_correctness=True
1.767 ms, 1555 TFLOP/s, name: fwd
8.214 ms, 870 TFLOP/s, name: bwd
b=2, mean_sq=4096, mean_sk=4096, varlen=False, h=128, h_k=128, d=192, dv=128, causal=True, has_bwd=True, check_correctness=True
1.086 ms, 1265 TFLOP/s, name: fwd
4.648 ms, 769 TFLOP/s, name: bwd
b=2, mean_sq=4096, mean_sk=4096, varlen=False, h=32, h_k=4, d=128, dv=128, causal=False, has_bwd=False, check_correctness=True
0.405 ms, 1356 TFLOP/s, name: fwd
b=2, mean_sq=4096, mean_sk=4096, varlen=False, h=32, h_k=4, d=128, dv=128, causal=True, has_bwd=False, check_correctness=True
0.253 ms, 1087 TFLOP/s, name: fwd
b=2, mean_sq=4096, mean_sk=4096, varlen=False, h=32, h_k=4, d=192, dv=128, causal=False, has_bwd=False, check_correctness=True
0.451 ms, 1525 TFLOP/s, name: fwd
b=2, mean_sq=4096, mean_sk=4096, varlen=False, h=32, h_k=4, d=192, dv=128, causal=True, has_bwd=False, check_correctness=True
0.283 ms, 1215 TFLOP/s, name: fwd
b=2, mean_sq=4096, mean_sk=4096, varlen=True, h=128, h_k=128, d=128, dv=128, causal=False, has_bwd=True, check_correctness=True
1.924 ms, 1290 TFLOP/s, name: fwd
6.322 ms, 982 TFLOP/s, name: bwd
b=2, mean_sq=4096, mean_sk=4096, varlen=True, h=128, h_k=128, d=128, dv=128, causal=True, has_bwd=True, check_correctness=True
1.324 ms, 1104 TFLOP/s, name: fwd
4.190 ms, 872 TFLOP/s, name: bwd
b=2, mean_sq=4096, mean_sk=4096, varlen=True, h=128, h_k=128, d=192, dv=128, causal=False, has_bwd=True, check_correctness=True
2.263 ms, 1371 TFLOP/s, name: fwd
9.370 ms, 861 TFLOP/s, name: bwd
b=2, mean_sq=4096, mean_sk=4096, varlen=True, h=128, h_k=128, d=192, dv=128, causal=True, has_bwd=True, check_correctness=True
1.470 ms, 1243 TFLOP/s, name: fwd
6.236 ms, 762 TFLOP/s, name: bwd
b=2, mean_sq=4096, mean_sk=4096, varlen=True, h=32, h_k=4, d=128, dv=128, causal=False, has_bwd=False, check_correctness=True
0.502 ms, 1236 TFLOP/s, name: fwd
b=2, mean_sq=4096, mean_sk=4096, varlen=True, h=32, h_k=4, d=128, dv=128, causal=True, has_bwd=False, check_correctness=True
0.344 ms, 1063 TFLOP/s, name: fwd
b=2, mean_sq=4096, mean_sk=4096, varlen=True, h=32, h_k=4, d=192, dv=128, causal=False, has_bwd=False, check_correctness=True
0.554 ms, 1400 TFLOP/s, name: fwd
b=2, mean_sq=4096, mean_sk=4096, varlen=True, h=32, h_k=4, d=192, dv=128, causal=True, has_bwd=False, check_correctness=True
0.377 ms, 1212 TFLOP/s, name: fwd
b=2, mean_sq=8192, mean_sk=8192, varlen=False, h=128, h_k=128, d=128, dv=128, causal=False, has_bwd=True, check_correctness=False
6.931 ms, 1269 TFLOP/s, name: fwd
20.880 ms, 1053 TFLOP/s, name: bwd
b=2, mean_sq=8192, mean_sk=8192, varlen=False, h=128, h_k=128, d=128, dv=128, causal=True, has_bwd=True, check_correctness=False
3.944 ms, 1115 TFLOP/s, name: fwd
12.433 ms, 884 TFLOP/s, name: bwd
b=2, mean_sq=8192, mean_sk=8192, varlen=False, h=128, h_k=128, d=192, dv=128, causal=False, has_bwd=True, check_correctness=False
8.139 ms, 1351 TFLOP/s, name: fwd
30.933 ms, 924 TFLOP/s, name: bwd
b=2, mean_sq=8192, mean_sk=8192, varlen=False, h=128, h_k=128, d=192, dv=128, causal=True, has_bwd=True, check_correctness=False
4.356 ms, 1262 TFLOP/s, name: fwd
16.764 ms, 853 TFLOP/s, name: bwd
b=2, mean_sq=8192, mean_sk=8192, varlen=False, h=32, h_k=4, d=128, dv=128, causal=False, has_bwd=False, check_correctness=True
1.548 ms, 1420 TFLOP/s, name: fwd
b=2, mean_sq=8192, mean_sk=8192, varlen=False, h=32, h_k=4, d=128, dv=128, causal=True, has_bwd=False, check_correctness=True
0.857 ms, 1283 TFLOP/s, name: fwd
b=2, mean_sq=8192, mean_sk=8192, varlen=False, h=32, h_k=4, d=192, dv=128, causal=False, has_bwd=False, check_correctness=True
1.716 ms, 1602 TFLOP/s, name: fwd
b=2, mean_sq=8192, mean_sk=8192, varlen=False, h=32, h_k=4, d=192, dv=128, causal=True, has_bwd=False, check_correctness=True
0.952 ms, 1444 TFLOP/s, name: fwd
b=2, mean_sq=8192, mean_sk=8192, varlen=True, h=128, h_k=128, d=128, dv=128, causal=False, has_bwd=True, check_correctness=False
8.344 ms, 1190 TFLOP/s, name: fwd
24.550 ms, 1011 TFLOP/s, name: bwd
b=2, mean_sq=8192, mean_sk=8192, varlen=True, h=128, h_k=128, d=128, dv=128, causal=True, has_bwd=True, check_correctness=False
5.496 ms, 1064 TFLOP/s, name: fwd
15.733 ms, 929 TFLOP/s, name: bwd
b=2, mean_sq=8192, mean_sk=8192, varlen=True, h=128, h_k=128, d=192, dv=128, causal=False, has_bwd=True, check_correctness=False
9.633 ms, 1289 TFLOP/s, name: fwd
36.049 ms, 895 TFLOP/s, name: bwd
b=2, mean_sq=8192, mean_sk=8192, varlen=True, h=128, h_k=128, d=192, dv=128, causal=True, has_bwd=True, check_correctness=False
6.381 ms, 1145 TFLOP/s, name: fwd
22.228 ms, 855 TFLOP/s, name: bwd
b=2, mean_sq=8192, mean_sk=8192, varlen=True, h=32, h_k=4, d=128, dv=128, causal=False, has_bwd=False, check_correctness=True
1.854 ms, 1339 TFLOP/s, name: fwd
b=2, mean_sq=8192, mean_sk=8192, varlen=True, h=32, h_k=4, d=128, dv=128, causal=True, has_bwd=False, check_correctness=True
1.150 ms, 1271 TFLOP/s, name: fwd
b=2, mean_sq=8192, mean_sk=8192, varlen=True, h=32, h_k=4, d=192, dv=128, causal=False, has_bwd=False, check_correctness=True
2.081 ms, 1491 TFLOP/s, name: fwd
b=2, mean_sq=8192, mean_sk=8192, varlen=True, h=32, h_k=4, d=192, dv=128, causal=True, has_bwd=False, check_correctness=True
1.275 ms, 1433 TFLOP/s, name: fwd

janeyx99 · 2026-07-02T16:39:37Z

-static inline PtrT* get_tensor_ptr(const at::Tensor& tensor) {
-    if (tensor.has_storage()) {
+static inline PtrT* get_tensor_ptr(const torch::stable::Tensor& tensor) {
+    if (tensor.defined()) {


This is a weaker check, but as these kernels are registered to the CUDA backend key, the inputs will either be CUDA storage-backed tensors or null tensors so this check is sufficient.

pybind->TORCH_LIBRARY; get ABI stable with CPython

4114849

janeyx99 force-pushed the abi-stable-migration branch 2 times, most recently from fa9c1ed to 3726f95 Compare July 2, 2026 16:25

Make FlashMLA ABI stable with torch

3911d00

janeyx99 force-pushed the abi-stable-migration branch from 3726f95 to 3911d00 Compare July 2, 2026 16:29

janeyx99 commented Jul 2, 2026

View reviewed changes

Add a README edit describing the fork and build instructions

1638811

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

pybind->TORCH_LIBRARY; get ABI stable with CPython#193

pybind->TORCH_LIBRARY; get ABI stable with CPython#193
janeyx99 wants to merge 3 commits into
deepseek-ai:mainfrom
janeyx99:abi-stable-migration

janeyx99 commented Jul 2, 2026 •

edited

Loading

Uh oh!

janeyx99 Jul 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

janeyx99 commented Jul 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Test Plan

Uh oh!

janeyx99 Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

janeyx99 commented Jul 2, 2026 •

edited

Loading