pybind->TORCH_LIBRARY; get ABI stable with CPython

README.md

@@ -1,5 +1,31 @@
 # FlashMLA
 
+This fork is a LibTorch and CPython ABI stable version of the original FlashMLA. This means that you can build one wheel that suffices across Python 3.10+ and torch 2.10+ at runtime with a `pip install`. You can build and install it directly from GitHub; the cutlass submodule will be fetched automatically during the build:
+
+```zsh
+pip install -v --no-build-isolation "git+https://github.com/janeyx99/FlashMLA-ABI-Stable.git"
+```
+
+**Build prerequisites**
+- PyTorch **2.11+**
+    - Example command for cuda 13.0: `pip install torch==2.11.0 --index-url https://download.pytorch.org/whl/cu130`
+    - Note that the build time requirement is intentionally 2.11 to access nicer UX features for the stable ABI. However at runtime, the wheel will work with torch 2.10+.
+    - `--no-build-isolation` is required so the build can import your existing torch. 
+- A CUDA toolkit with `nvcc` **12.9+** for the SM100 (Blackwell) kernels. With `nvcc` 12.8 you must disable SM100: prefix the command with `FLASH_MLA_DISABLE_SM100=1`. 
+
+### Install from a local clone
+
+```zsh
+git clone https://github.com/janeyx99/FlashMLA-ABI-Stable.git
+cd FlashMLA-ABI-Stable
+git submodule update --init --recursive
+pip install -v --no-build-isolation .
+```
+
+Below is the original README with all the deets -- credits to the original authors from DeepSeek.
+
+
+
 ## Introduction
 
 FlashMLA is DeepSeek's library of optimized attention kernels, powering the [DeepSeek-V3](https://github.com/deepseek-ai/DeepSeek-V3) and [DeepSeek-V3.2-Exp](https://github.com/deepseek-ai/DeepSeek-V3.2-Exp) models. This repository contains the following implementations:

csrc/api/api.cpp

@@ -1,15 +1,22 @@
-#include <pybind11/pybind11.h>
+#include <torch/csrc/stable/library.h>
 
 #include "sparse_fwd.h"
 #include "sparse_decode.h"
 #include "dense_decode.h"
 #include "dense_fwd.h"
 
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.doc() = "FlashMLA";
-    m.def("sparse_decode_fwd", &sparse_attn_decode_interface);
-    m.def("dense_decode_fwd", &dense_attn_decode_interface);
-    m.def("sparse_prefill_fwd", &sparse_attn_prefill_interface);
-    m.def("dense_prefill_fwd", &FMHACutlassSM100FwdRun);
-    m.def("dense_prefill_bwd", &FMHACutlassSM100BwdRun);
+STABLE_TORCH_LIBRARY(flash_mla, m) {
+    m.def("sparse_decode_fwd(Tensor q, Tensor kv, Tensor indices, Tensor? topk_length, Tensor? attn_sink, Tensor(a)? tile_scheduler_metadata, Tensor(b)? num_splits, Tensor? extra_kv, Tensor? extra_indices, Tensor? extra_topk_length, int d_v, float sm_scale) -> (Tensor, Tensor, Tensor(a)?, Tensor(b)?)");
+    m.def("dense_decode_fwd(Tensor q, Tensor kcache, int head_size_v, Tensor seqlens_k, Tensor block_table, float softmax_scale, bool is_causal, Tensor(a)? tile_scheduler_metadata, Tensor(b)? num_splits) -> (Tensor, Tensor, Tensor(a)?, Tensor(b)?)");
+    m.def("sparse_prefill_fwd(Tensor q, Tensor kv, Tensor indices, float sm_scale, int d_v, Tensor? attn_sink, Tensor? topk_length) -> Tensor[]");
+    m.def("dense_prefill_fwd(Tensor workspace_buffer, Tensor q, Tensor k, Tensor v, Tensor cumulative_seqlen_q, Tensor cumulative_seqlen_kv, Tensor(a!) o, Tensor(b!) lse, int mask_mode_code, float softmax_scale, int max_seqlen_q, int max_seqlen_kv, bool is_varlen) -> ()");
+    m.def("dense_prefill_bwd(Tensor(a!) workspace_buffer, Tensor d_o, Tensor q, Tensor k, Tensor v, Tensor o, Tensor lse, Tensor cumulative_seqlen_q, Tensor cumulative_seqlen_kv, Tensor(b!) dq, Tensor(c!) dk, Tensor(d!) dv, int mask_mode_code, float softmax_scale, int max_seqlen_q, int max_seqlen_kv, bool is_varlen) -> ()");
+}
+
+STABLE_TORCH_LIBRARY_IMPL(flash_mla, CUDA, m) {
+    m.impl("sparse_decode_fwd", TORCH_BOX(&sparse_attn_decode_interface));
+    m.impl("dense_decode_fwd", TORCH_BOX(&dense_attn_decode_interface));
+    m.impl("sparse_prefill_fwd", TORCH_BOX(&sparse_attn_prefill_interface));
+    m.impl("dense_prefill_fwd", TORCH_BOX(&FMHACutlassSM100FwdRun));
+    m.impl("dense_prefill_bwd", TORCH_BOX(&FMHACutlassSM100BwdRun));
 }

csrc/api/common.h

@@ -1,31 +1,62 @@
 #pragma once
 
 #include <span>
+#include <array>
+
+#include <cuda_runtime.h>
+
+#include <torch/csrc/stable/tensor.h>
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/ops.h>
+#include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+
+#include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/util/Exception.h>
+#include <torch/headeronly/util/shim_utils.h>
 
-#include <torch/extension.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
 #include <kerutils/supplemental/torch_tensors.h>
+#include <kerutils/supplemental/cuda_stream.h>
 
 #include <cutlass/bfloat16.h>
 
+using torch::stable::Tensor;
+using torch::headeronly::ScalarType;
+
+// Re-exported from kerutils so existing call sites can use it unqualified.
+using kerutils::get_current_cuda_stream;
+
 static constexpr float LOG_2_E = 1.44269504f;
 
-// Instantiation for tensor.data_ptr<cutlass::bfloat16_t>()
-template<>
-inline cutlass::bfloat16_t* at::TensorBase::data_ptr<cutlass::bfloat16_t>() const {
-    return reinterpret_cast<cutlass::bfloat16_t*>(this->data_ptr());
+// Helper to access device properties (SM count, compute capability) via the CUDA
+// Runtime directly and cache per device index.
+inline const cudaDeviceProp &get_cached_device_prop() {
+    int device_idx = static_cast<int>(torch::stable::accelerator::getCurrentDeviceIndex());
+    constexpr int kMaxDevices = 16;
+    static cudaDeviceProp props[kMaxDevices];
+    static bool inited[kMaxDevices] = {false};
+    if (device_idx < 0 || device_idx >= kMaxDevices) {
+        // Uncached fallback for unexpectedly large device indices.
+        static thread_local cudaDeviceProp tmp;
+        cudaGetDeviceProperties(&tmp, device_idx);
+        return tmp;
+    }
+    if (!inited[device_idx]) {
+        cudaGetDeviceProperties(&props[device_idx], device_idx);
+        inited[device_idx] = true;
+    }
+    return props[device_idx];
 }
 
 // A struct that holds the architecture information of the current GPU.
 struct Arch {
     int major;
     int minor;
     int num_sms;
-    cudaDeviceProp* device_prop;
+    const cudaDeviceProp* device_prop;
 
     Arch() {
-        device_prop = at::cuda::getCurrentDeviceProperties();
+        device_prop = &get_cached_device_prop();
         major = device_prop->major;
         minor = device_prop->minor;
         num_sms = device_prop->multiProcessorCount;
@@ -43,7 +74,7 @@ struct Arch {
 // Convert int64_t stride to int32_t, with overflow check.
 inline int int64_stride_to_int(int64_t orig_stride) {
     if (orig_stride > std::numeric_limits<int>::max()) {
-        TORCH_CHECK(false, "[FlashMLA] Stride exceeds int32 limit: ", orig_stride);
+        STD_TORCH_CHECK(false, "[FlashMLA] Stride exceeds int32 limit: ", orig_stride);
     }
     return static_cast<int>(orig_stride);
 }
@@ -57,7 +88,7 @@ inline int int64_stride_to_int(int64_t orig_stride) {
             static constexpr int CONSTEXPR_NAME = 64; \
             return __VA_ARGS__(); \
         } else { \
-            TORCH_CHECK(false, "Unsupported num_heads_q: ", NUM_HEADS); \
+            STD_TORCH_CHECK(false, "Unsupported num_heads_q: ", NUM_HEADS); \
         } \
     } ();
 
@@ -70,7 +101,7 @@ inline int int64_stride_to_int(int64_t orig_stride) {
         static constexpr int CONSTEXPR_NAME = 512; \
         return __VA_ARGS__(); \
     } else { \
-        TORCH_CHECK(false, "Unsupported head_dim_qk: ", HEAD_DIM); \
+        STD_TORCH_CHECK(false, "Unsupported head_dim_qk: ", HEAD_DIM); \
     } \
 } ();
 
@@ -94,7 +125,7 @@ inline int int64_stride_to_int(int64_t orig_stride) {
         static constexpr ModelType CONSTEXPR_NAME = ModelType::MODEL1; \
         return __VA_ARGS__(); \
     } else { \
-        TORCH_CHECK(false, "Unsupported model type: ", (int)MODEL_TYPE); \
+        STD_TORCH_CHECK(false, "Unsupported model type: ", (int)MODEL_TYPE); \
     } \
 } ();
 
@@ -120,7 +151,7 @@ constexpr auto get_static_enum_name(){
     };
 }
 
-template<typename T, std::size_t N = 0> 
+template<typename T, std::size_t N = 0>
 static constexpr std::size_t get_enum_max(){
     constexpr T value = static_cast<T>(N);
     if constexpr (get_static_enum_name<value>().find(")") == std::string_view::npos)
@@ -133,8 +164,8 @@ template<typename T> requires std::is_enum_v<T>
 static constexpr std::string get_dynamic_enum_name(T value){
     constexpr std::size_t num = get_enum_max<T>();
     constexpr auto names = []<std::size_t... Is>(std::index_sequence<Is...>){
-        return std::array<std::string_view, num>{ 
-            get_static_enum_name<static_cast<T>(Is)>()... 
+        return std::array<std::string_view, num>{
+            get_static_enum_name<static_cast<T>(Is)>()...
         };
     }(std::make_index_sequence<num>{});
     return (std::string)names[static_cast<std::size_t>(value)];
@@ -219,7 +250,7 @@ class ImplBase {
             Arch cur_gpu_arch = Arch();
             fprintf(stderr, "Current GPU: %s, SM %d.%d with %d SMs\n", cur_gpu_arch.device_prop->name, cur_gpu_arch.major, cur_gpu_arch.minor, cur_gpu_arch.num_sms);
             fprintf(stderr, "This means that the dispatcher has chosen an implementation that does not support all required features. Maybe there is a bug in the dispatcher, or you have requested an invalid combination of features.\n");
-            TORCH_CHECK(false, "The chosen implementation does not support all required features. See message above for details.");
+            STD_TORCH_CHECK(false, "The chosen implementation does not support all required features. See message above for details.");
         }
     }
 
@@ -228,4 +259,3 @@ class ImplBase {
         run_(params, required_features);
     }
 };
-