feat(moore): call MATE flash attention natively

src/CMakeLists.txt

@@ -843,6 +843,35 @@ if(_infini_ops_linked_uses_torch)
     list(APPEND TORCH_SOURCES ${INFINI_OPS_LINKED_TORCH_SOURCES})
 endif()
 
+if(WITH_MOORE AND WITH_LINKED)
+    set(_moore_mate_linked_ops
+        flash_attn_varlen_func
+        flash_attn_with_kvcache)
+    set(_moore_mate_linked_enabled FALSE)
+    set(_moore_mate_sources "")
+    foreach(_op IN LISTS _moore_mate_linked_ops)
+        set(_moore_mate_source
+            "${CMAKE_CURRENT_SOURCE_DIR}/linked/tvm_ffi/moore/ops/${_op}/mate.cc")
+        list(FIND INFINI_OPS_LINKED_SOURCES "${_moore_mate_source}"
+            _moore_mate_source_index)
+        if(NOT _moore_mate_source_index EQUAL -1)
+            list(APPEND _moore_mate_sources "${_moore_mate_source}")
+            set(_moore_mate_linked_enabled TRUE)
+        endif()
+    endforeach()
+
+    if(_moore_mate_linked_enabled)
+        find_package(Python COMPONENTS Interpreter Development REQUIRED)
+        find_library(TORCH_PYTHON_LIB torch_python
+            HINTS ${_torch_lib_dirs} REQUIRED)
+        list(APPEND TORCH_INCLUDE_DIRS ${Python_INCLUDE_DIRS})
+        list(APPEND TORCH_LIBRARIES
+            ${TORCH_PYTHON_LIB} Python::Python ${CMAKE_DL_LIBS})
+        set(_infini_ops_linked_uses_torch TRUE)
+        list(APPEND TORCH_SOURCES ${_moore_mate_sources})
+        list(REMOVE_ITEM INFINI_OPS_LINKED_TVM_FFI_SOURCES ${_moore_mate_sources})
+    endif()
+endif()
 
 if(_infini_ops_linked_uses_tvm_ffi)
     target_sources(infiniops PRIVATE ${INFINI_OPS_LINKED_TVM_FFI_SOURCES})
@@ -929,7 +958,13 @@ endif()
 if(TORCH_SOURCES)
     set(INFINI_OPS_TORCH_UNITY_BATCH_SIZE "8" CACHE STRING
         "Number of torch sources to include in each generated unity translation unit; set to 1 to disable")
-    set(TORCH_COMPILE_SOURCES ${TORCH_SOURCES})
+    set(_torch_batchable_sources ${TORCH_SOURCES})
+    set(_torch_unbatched_sources "")
+    if(_moore_mate_linked_enabled)
+        list(REMOVE_ITEM _torch_batchable_sources ${_moore_mate_sources})
+        list(APPEND _torch_unbatched_sources ${_moore_mate_sources})
+    endif()
+    set(TORCH_COMPILE_SOURCES ${_torch_batchable_sources})
     if(INFINI_OPS_TORCH_UNITY_BATCH_SIZE GREATER 1)
         set(_torch_unity_dir "${CMAKE_CURRENT_BINARY_DIR}/torch_unity")
         file(MAKE_DIRECTORY "${_torch_unity_dir}")
@@ -939,7 +974,7 @@ if(TORCH_SOURCES)
         set(_torch_unity_index 0)
         set(_torch_unity_count 0)
         set(_torch_unity_content "")
-        foreach(_src IN LISTS TORCH_SOURCES)
+        foreach(_src IN LISTS _torch_batchable_sources)
             if(_torch_unity_count EQUAL 0)
                 set(_torch_unity_src
                     "${_torch_unity_dir}/torch_unity_${_torch_unity_index}.cc")
@@ -974,6 +1009,10 @@ if(TORCH_SOURCES)
             "${_torch_unity_source_count} translation units")
     endif()
 
+    if(_torch_unbatched_sources)
+        list(APPEND TORCH_COMPILE_SOURCES ${_torch_unbatched_sources})
+    endif()
+
     if(WITH_TORCH)
         target_compile_definitions(infiniops PUBLIC
             $<BUILD_INTERFACE:WITH_TORCH=1>)
@@ -1010,7 +1049,7 @@ if(TORCH_SOURCES)
         endif()
 
         set(_torch_include_flags "")
-        foreach(_dir ${TORCH_INCLUDE_DIRS})
+        foreach(_dir ${TORCH_INCLUDE_DIRS} ${INFINI_OPS_LINKED_INCLUDE_DIRS})
             list(APPEND _torch_include_flags "-isystem" "${_dir}")
         endforeach()
 

src/linked/tvm_ffi/moore/mate.h

@@ -0,0 +1,227 @@
+#ifndef INFINI_OPS_LINKED_TVM_FFI_MOORE_MATE_H_
+#define INFINI_OPS_LINKED_TVM_FFI_MOORE_MATE_H_
+
+#include <ATen/DLConvertor.h>
+#include <dlfcn.h>
+#include <dlpack/dlpack.h>
+#include <pybind11/pybind11.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/extra/c_env_api.h>
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/optional.h>
+
+#include <cstdio>
+#include <fstream>
+#include <memory>
+#include <optional>
+#include <stdexcept>
+#include <string>
+#include <utility>
+
+namespace infini::ops::linked::tvm_ffi::moore {
+
+namespace py = pybind11;
+
+using TvmFfiEntry = int(void*, const TVMFFIAny*, int32_t, TVMFFIAny*);
+using OptionalTensorView = tvm::ffi::Optional<tvm::ffi::TensorView>;
+
+namespace detail {
+
+struct DlLibraryCloser {
+  void operator()(void* library) const {
+    if (library != nullptr && dlclose(library) != 0) {
+      std::fprintf(stderr, "[InfiniOps] MATE library close: %s\n", dlerror());
+    }
+  }
+};
+
+using DlLibrary = std::unique_ptr<void, DlLibraryCloser>;
+
+class DlSymbol {
+ public:
+  void Load(const std::string& dispatch_name) {
+    const auto library_path = FindLoadedLibrary(dispatch_name);
+    library_ = DlLibrary{dlopen(library_path.c_str(), RTLD_NOW | RTLD_LOCAL)};
+    if (library_ == nullptr) {
+      throw std::runtime_error("MATE failed to load " + library_path + ": " +
+                               dlerror());
+    }
+
+    dlerror();
+    entry_ = reinterpret_cast<TvmFfiEntry*>(
+        dlsym(library_.get(), ("__tvm_ffi_" + dispatch_name).c_str()));
+    if (const auto* error = dlerror(); entry_ == nullptr || error != nullptr) {
+      throw std::runtime_error("MATE did not export __tvm_ffi_" +
+                               dispatch_name + ": " + error);
+    }
+  }
+
+  TvmFfiEntry* Entry() const { return entry_; }
+
+ private:
+  static std::string FindLoadedLibrary(const std::string& dispatch_name) {
+    std::ifstream maps("/proc/self/maps");
+    std::string line;
+    const auto directory = "/" + dispatch_name + "/";
+    const auto basename = dispatch_name + ".so";
+    while (std::getline(maps, line)) {
+      const auto path_pos = line.find(" /");
+      if (path_pos == std::string::npos) continue;
+      auto path = line.substr(path_pos + 1);
+      if (path.find(directory) != std::string::npos ||
+          path.substr(path.find_last_of('/') + 1) == basename) {
+        return path;
+      }
+    }
+    throw std::runtime_error("MATE did not retain the loaded module for " +
+                             dispatch_name);
+  }
+
+  DlLibrary library_;
+  TvmFfiEntry* entry_{nullptr};
+};
+
+class DlPackTensor {
+ public:
+  explicit DlPackTensor(const at::Tensor& tensor)
+      : managed_(at::toDLPack(tensor)) {}
+
+  DlPackTensor(const DlPackTensor&) = delete;
+  DlPackTensor& operator=(const DlPackTensor&) = delete;
+
+  ~DlPackTensor() {
+    if (managed_ != nullptr && managed_->deleter != nullptr) {
+      managed_->deleter(managed_);
+    }
+  }
+
+  const DLTensor* Get() const { return &managed_->dl_tensor; }
+
+ private:
+  DLManagedTensor* managed_;
+};
+
+class TvmStreamGuard {
+ public:
+  TvmStreamGuard(DLDevice device, void* stream) : device_{device} {
+    const auto status = TVMFFIEnvSetStream(device.device_type, device.device_id,
+                                           stream, &previous_stream_);
+    if (status != 0) {
+      throw std::runtime_error(
+          "MATE failed to select the TVM-FFI MUSA stream (status " +
+          std::to_string(status) + ")");
+    }
+  }
+
+  ~TvmStreamGuard() {
+    const auto status = TVMFFIEnvSetStream(
+        device_.device_type, device_.device_id, previous_stream_, nullptr);
+    if (status != 0) {
+      std::fprintf(stderr,
+                   "[InfiniOps] MATE failed to restore the TVM-FFI stream "
+                   "(status %d)\n",
+                   status);
+    }
+  }
+
+ private:
+  DLDevice device_;
+  TVMFFIStreamHandle previous_stream_{nullptr};
+};
+
+class ModuleRecorder {
+ public:
+  ModuleRecorder() {
+    const auto mate = py::module_::import("mate");
+    const auto version = py::str(mate.attr("__version__")).cast<std::string>();
+    const auto separator = version.find('+');
+    if (version.substr(0, separator) != "0.2.5") {
+      throw std::runtime_error(
+          "Mate 0.2.5 is required by the Moore native "
+          "FlashAttention provider, but found " +
+          version);
+    }
+
+    forward_ = py::module_::import("mate.jit.attention.fmha.fmha_fwd");
+    combine_ = py::module_::import("mate.jit.attention.fmha.fmha_combine");
+    original_forward_loader_ = forward_.attr("_fmha_fwd_module");
+    original_combine_loader_ = combine_.attr("_fmha_fwd_combine_module");
+    forward_names_ = py::list();
+    combine_names_ = py::list();
+
+    auto forward_loader = py::cpp_function(
+        [forward = forward_, original = original_forward_loader_,
+         names = forward_names_](py::object config) {
+          names.append(forward.attr("_fmha_fwd_encode")(config));
+          return original(config);
+        });
+    auto combine_loader = py::cpp_function(
+        [combine = combine_, original = original_combine_loader_,
+         names = combine_names_](py::object config) {
+          names.append(combine.attr("_fmha_fwd_combine_encode")(config));
+          return original(config);
+        });
+    forward_.attr("_fmha_fwd_module") = forward_loader;
+    combine_.attr("_fmha_fwd_combine_module") = combine_loader;
+  }
+
+  ModuleRecorder(const ModuleRecorder&) = delete;
+  ModuleRecorder& operator=(const ModuleRecorder&) = delete;
+
+  ~ModuleRecorder() {
+    if (forward_) forward_.attr("_fmha_fwd_module") = original_forward_loader_;
+    if (combine_) {
+      combine_.attr("_fmha_fwd_combine_module") = original_combine_loader_;
+    }
+  }
+
+  py::module_ Forward() const { return forward_; }
+
+  std::string ForwardName() const { return Name(forward_names_); }
+
+  std::optional<std::string> CombineName() const {
+    if (py::len(combine_names_) == 0) return std::nullopt;
+    return Name(combine_names_);
+  }
+
+ private:
+  static std::string Name(const py::list& names) {
+    if (py::len(names) != 1) {
+      throw std::runtime_error(
+          "MATE selected an unexpected number of FlashAttention modules");
+    }
+    return py::str(names[0]).cast<std::string>();
+  }
+
+  py::module_ forward_;
+  py::module_ combine_;
+  py::object original_forward_loader_;
+  py::object original_combine_loader_;
+  py::list forward_names_;
+  py::list combine_names_;
+};
+
+}  // namespace detail
+
+class MateFmhaRuntime {
+ public:
+  void Load(const std::string& forward_name,
+            const std::optional<std::string>& combine_name) {
+    forward_.Load(forward_name);
+    if (combine_name.has_value()) combine_.Load(*combine_name);
+  }
+
+  detail::DlSymbol& Forward() { return forward_; }
+
+  detail::DlSymbol* Combine() {
+    return combine_.Entry() == nullptr ? nullptr : &combine_;
+  }
+
+ private:
+  detail::DlSymbol forward_;
+  detail::DlSymbol combine_;
+};
+
+}  // namespace infini::ops::linked::tvm_ffi::moore
+
+#endif  // INFINI_OPS_LINKED_TVM_FFI_MOORE_MATE_H_