diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index 4e475f3c6..24ebad24f 100644
--- a/src/CMakeLists.txt
+++ b/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()
 
diff --git a/src/linked/tvm_ffi/moore/mate.h b/src/linked/tvm_ffi/moore/mate.h
new file mode 100644
index 000000000..1ad96bab9
--- /dev/null
+++ b/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_
diff --git a/src/linked/tvm_ffi/moore/ops/flash_attn_varlen_func/mate.cc b/src/linked/tvm_ffi/moore/ops/flash_attn_varlen_func/mate.cc
new file mode 100644
index 000000000..a611adad8
--- /dev/null
+++ b/src/linked/tvm_ffi/moore/ops/flash_attn_varlen_func/mate.cc
@@ -0,0 +1,222 @@
+#include "linked/tvm_ffi/moore/ops/flash_attn_varlen_func/mate.h"
+
+#include <tvm/ffi/container/array.h>
+
+#include <cmath>
+#include <stdexcept>
+#include <vector>
+
+#include "torch/moore/c10.h"
+#include "torch/tensor_.h"
+
+namespace infini::ops {
+namespace {
+
+namespace mate = linked::tvm_ffi::moore;
+using mate::OptionalTensorView;
+
+py::object PythonTensor(const std::optional<at::Tensor>& tensor) {
+  return tensor.has_value() ? py::cast(*tensor) : py::none();
+}
+
+void CallCombine(mate::MateFmhaRuntime& runtime,
+                 mate::detail::DlPackTensor& out,
+                 mate::detail::DlPackTensor& lse,
+                 const tvm::ffi::Array<tvm::ffi::Tensor>& accumulators,
+                 mate::detail::DlPackTensor& cu_seqlens_q, int max_seqlen_q,
+                 int num_splits) {
+  auto* combine = runtime.Combine();
+  if (num_splits > 1 && combine == nullptr) {
+    throw std::runtime_error(
+        "MATE selected split-KV execution without a combine kernel");
+  }
+  if (num_splits <= 1) return;
+
+  tvm::ffi::Function::InvokeExternC(
+      nullptr, combine->Entry(),
+      OptionalTensorView{tvm::ffi::TensorView(cu_seqlens_q.Get())},
+      OptionalTensorView{}, tvm::ffi::Optional<int>{max_seqlen_q},
+      tvm::ffi::TensorView(out.Get()), tvm::ffi::TensorView(lse.Get()),
+      tvm::ffi::TensorView(accumulators[0]),
+      tvm::ffi::TensorView(accumulators[1]), OptionalTensorView{},
+      int{num_splits});
+}
+
+}  // namespace
+
+void Operator<FlashAttnVarlenFunc, Device::Type::kMoore, 16>::operator()(
+    const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q,
+    const Tensor cu_seqlens_k, const std::optional<Tensor> alibi_slopes,
+    const std::optional<Tensor> block_table, const int64_t max_seqlen_q,
+    const int64_t max_seqlen_k, const double dropout_p,
+    const std::optional<double> softmax_scale, const bool causal,
+    const std::vector<int64_t> window_size, const double softcap,
+    const bool deterministic, const bool return_attn_probs, Tensor out,
+    std::optional<Tensor> softmax_lse, std::optional<Tensor> s_dmask) const {
+  std::lock_guard lock{runtime_mutex_};
+  Call(q, k, v, cu_seqlens_q, cu_seqlens_k, alibi_slopes, block_table,
+       max_seqlen_q, max_seqlen_k, dropout_p, softmax_scale, causal,
+       window_size, softcap, deterministic, return_attn_probs, out, softmax_lse,
+       s_dmask);
+}
+
+void Operator<FlashAttnVarlenFunc, Device::Type::kMoore, 16>::Call(
+    const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q,
+    const Tensor cu_seqlens_k, const std::optional<Tensor> alibi_slopes,
+    const std::optional<Tensor> block_table, const int64_t max_seqlen_q,
+    const int64_t max_seqlen_k, const double dropout_p,
+    const std::optional<double> softmax_scale, const bool causal,
+    const std::vector<int64_t> window_size, const double softcap,
+    const bool deterministic, const bool return_attn_probs, Tensor out,
+    std::optional<Tensor> softmax_lse, std::optional<Tensor> s_dmask) const {
+  (void)s_dmask;
+  TORCH_CHECK(!alibi_slopes.has_value(), "MATE ALiBi is not supported");
+  TORCH_CHECK(dropout_p == 0.0, "MATE attention dropout is not supported");
+  TORCH_CHECK(!deterministic, "MATE deterministic attention is not supported");
+  TORCH_CHECK(!return_attn_probs,
+              "MATE attention probabilities are not supported");
+
+  const typename C10<Device::Type::kMoore>::StreamGuard stream_guard{
+      C10<Device::Type::kMoore>::GetStreamFromExternal(stream_, device_index_)};
+
+  auto at_q =
+      ToAtenTensor<Device::Type::kMoore>(const_cast<void*>(q.data()), q_shape_,
+                                         q_strides_, q_dtype_, device_index_);
+  auto at_k =
+      ToAtenTensor<Device::Type::kMoore>(const_cast<void*>(k.data()), k_shape_,
+                                         k_strides_, k_dtype_, device_index_);
+  auto at_v =
+      ToAtenTensor<Device::Type::kMoore>(const_cast<void*>(v.data()), v_shape_,
+                                         v_strides_, v_dtype_, device_index_);
+  auto at_cu_seqlens_q = ToAtenTensor<Device::Type::kMoore>(
+      const_cast<void*>(cu_seqlens_q.data()), cu_seqlens_q_shape_,
+      cu_seqlens_q_strides_, cu_seqlens_q_dtype_, device_index_);
+  auto at_cu_seqlens_k = ToAtenTensor<Device::Type::kMoore>(
+      const_cast<void*>(cu_seqlens_k.data()), cu_seqlens_k_shape_,
+      cu_seqlens_k_strides_, cu_seqlens_k_dtype_, device_index_);
+  auto at_out = ToAtenTensor<Device::Type::kMoore>(
+      out.data(), out_shape_, out_strides_, out_dtype_, device_index_);
+
+  std::optional<at::Tensor> at_block_table;
+  if (block_table.has_value()) {
+    at_block_table.emplace(ToAtenTensor<Device::Type::kMoore>(
+        const_cast<void*>(block_table->data()), block_table_shape_,
+        block_table_strides_, block_table_dtype_, device_index_));
+  }
+
+  if (block_table.has_value() && !paged_seqused_k_.has_value()) {
+    paged_seqused_k_.emplace(
+        (at_cu_seqlens_k.slice(0, 1) - at_cu_seqlens_k.slice(0, 0, -1))
+            .contiguous());
+  }
+  if (!softmax_lse.has_value() && !internal_lse_.has_value()) {
+    internal_lse_.emplace(at::empty(
+        {static_cast<int64_t>(q_shape_[1]), static_cast<int64_t>(q_shape_[0])},
+        at_out.options().dtype(at::kFloat)));
+  }
+  const auto& at_lse = softmax_lse.has_value() ? *softmax_lse : *internal_lse_;
+
+  if (runtime_.Forward().Entry() == nullptr) {
+    py::gil_scoped_acquire gil;
+    try {
+      mate::detail::ModuleRecorder recorder;
+      recorder.Forward().attr("_fmha_fwd")(
+          py::arg("q") = at_q, py::arg("k") = at_k, py::arg("v") = at_v,
+          py::arg("k_new") = py::none(), py::arg("v_new") = py::none(),
+          py::arg("q_v") = py::none(),
+          py::arg("cu_seqlens_q") = at_cu_seqlens_q,
+          py::arg("cu_seqlens_k") =
+              (block_table.has_value() ? py::none()
+                                       : py::cast(at_cu_seqlens_k)),
+          py::arg("cu_seqlens_k_new") = py::none(),
+          py::arg("seqused_q") = py::none(),
+          py::arg("seqused_k") = PythonTensor(paged_seqused_k_),
+          py::arg("max_seqlen_q") = max_seqlen_q,
+          py::arg("max_seqlen_k") = max_seqlen_k,
+          py::arg("page_table") = PythonTensor(at_block_table),
+          py::arg("kv_batch_idx") = py::none(),
+          py::arg("leftpad_k") = py::none(), py::arg("rotary_cos") = py::none(),
+          py::arg("rotary_sin") = py::none(),
+          py::arg("seqlens_rotary") = py::none(),
+          py::arg("q_descale") = py::none(), py::arg("k_descale") = py::none(),
+          py::arg("v_descale") = py::none(),
+          py::arg("softmax_scale") = static_cast<float>(softmax_scale.value_or(
+              1.0 / std::sqrt(static_cast<double>(q_shape_[2])))),
+          py::arg("is_causal") = causal,
+          py::arg("window_size_left") = window_size[0],
+          py::arg("window_size_right") = window_size[1],
+          py::arg("attention_chunk") = 0,
+          py::arg("learnable_sink") = py::none(),
+          py::arg("softcap") = static_cast<float>(softcap),
+          py::arg("is_rotary_interleaved") = false,
+          py::arg("scheduler_metadata") = py::none(), py::arg("num_splits") = 0,
+          py::arg("pack_gqa") =
+              (q_shape_[1] !=
+               (block_table.has_value() ? k_shape_[3] : k_shape_[2])),
+          py::arg("mp_margin") = 0, py::arg("return_lse") = true,
+          py::arg("lse") = at_lse, py::arg("out") = at_out,
+          py::arg("cp_world_size") = 1, py::arg("cp_rank") = 0,
+          py::arg("cp_tot_seqused_k") = py::none(), py::arg("only_qv") = false);
+      runtime_.Load(recorder.ForwardName(), recorder.CombineName());
+    } catch (const py::error_already_set& error) {
+      TORCH_CHECK(false, "MATE flash_attn_varlen_func bootstrap failed: ",
+                  error.what());
+    }
+    return;
+  }
+
+  mate::detail::DlPackTensor dl_q{at_q};
+  mate::detail::DlPackTensor dl_k{at_k};
+  mate::detail::DlPackTensor dl_v{at_v};
+  mate::detail::DlPackTensor dl_cu_q{at_cu_seqlens_q};
+  mate::detail::DlPackTensor dl_cu_k{at_cu_seqlens_k};
+  mate::detail::DlPackTensor dl_out{at_out};
+  mate::detail::DlPackTensor dl_lse{at_lse};
+  std::optional<mate::detail::DlPackTensor> dl_block_table;
+  std::optional<mate::detail::DlPackTensor> dl_seqused_k;
+  if (at_block_table.has_value()) {
+    dl_block_table.emplace(*at_block_table);
+    dl_seqused_k.emplace(*paged_seqused_k_);
+  }
+
+  const mate::detail::TvmStreamGuard tvm_stream_guard{
+      dl_q.Get()->device,
+      C10<Device::Type::kMoore>::GetStreamFromExternal(stream_, device_index_)};
+
+  auto result = tvm::ffi::Function::InvokeExternC(
+      nullptr, runtime_.Forward().Entry(), tvm::ffi::TensorView(dl_q.Get()),
+      tvm::ffi::TensorView(dl_k.Get()), tvm::ffi::TensorView(dl_v.Get()),
+      OptionalTensorView{}, OptionalTensorView{}, OptionalTensorView{},
+      tvm::ffi::TensorView(dl_cu_q.Get()),
+      block_table.has_value()
+          ? OptionalTensorView{}
+          : OptionalTensorView{tvm::ffi::TensorView(dl_cu_k.Get())},
+      OptionalTensorView{}, OptionalTensorView{},
+      dl_seqused_k.has_value()
+          ? OptionalTensorView{tvm::ffi::TensorView(dl_seqused_k->Get())}
+          : OptionalTensorView{},
+      tvm::ffi::Optional<int>{static_cast<int>(max_seqlen_q)},
+      tvm::ffi::Optional<int>{static_cast<int>(max_seqlen_k)},
+      dl_block_table.has_value()
+          ? OptionalTensorView{tvm::ffi::TensorView(dl_block_table->Get())}
+          : OptionalTensorView{},
+      OptionalTensorView{}, OptionalTensorView{}, OptionalTensorView{},
+      OptionalTensorView{}, OptionalTensorView{}, OptionalTensorView{},
+      OptionalTensorView{}, OptionalTensorView{}, OptionalTensorView{},
+      static_cast<float>(softmax_scale.value_or(
+          1.0 / std::sqrt(static_cast<double>(q_shape_[2])))),
+      causal, int{window_size[0]}, int{window_size[1]}, int{0},
+      static_cast<float>(softcap), int{0}, int{0}, OptionalTensorView{},
+      OptionalTensorView{}, tvm::ffi::TensorView(dl_out.Get()),
+      tvm::ffi::TensorView(dl_lse.Get()), int{1}, int{0}, OptionalTensorView{},
+      false);
+
+  const auto outputs = result.cast<tvm::ffi::Array<tvm::ffi::Any>>();
+  const auto accumulators =
+      outputs[0].cast<tvm::ffi::Array<tvm::ffi::Tensor>>();
+  const auto num_splits = outputs[1].cast<int>();
+  CallCombine(runtime_, dl_out, dl_lse, accumulators, dl_cu_q,
+              static_cast<int>(max_seqlen_q), num_splits);
+}
+
+}  // namespace infini::ops
diff --git a/src/linked/tvm_ffi/moore/ops/flash_attn_varlen_func/mate.h b/src/linked/tvm_ffi/moore/ops/flash_attn_varlen_func/mate.h
new file mode 100644
index 000000000..5b1200889
--- /dev/null
+++ b/src/linked/tvm_ffi/moore/ops/flash_attn_varlen_func/mate.h
@@ -0,0 +1,53 @@
+#ifndef INFINI_OPS_LINKED_TVM_FFI_MOORE_OPS_FLASH_ATTN_VARLEN_FUNC_MATE_H_
+#define INFINI_OPS_LINKED_TVM_FFI_MOORE_OPS_FLASH_ATTN_VARLEN_FUNC_MATE_H_
+
+#include <ATen/core/Tensor.h>
+
+#include <mutex>
+#include <optional>
+
+#include "base/flash_attn_varlen_func.h"
+#include "linked/tvm_ffi/moore/mate.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<FlashAttnVarlenFunc, Device::Type::kMoore, 16>
+    : public FlashAttnVarlenFunc {
+ public:
+  using FlashAttnVarlenFunc::FlashAttnVarlenFunc;
+  using FlashAttnVarlenFunc::operator();
+
+  void operator()(const Tensor q, const Tensor k, const Tensor v,
+                  const Tensor cu_seqlens_q, const Tensor cu_seqlens_k,
+                  const std::optional<Tensor> alibi_slopes,
+                  const std::optional<Tensor> block_table,
+                  const int64_t max_seqlen_q, const int64_t max_seqlen_k,
+                  const double dropout_p,
+                  const std::optional<double> softmax_scale, const bool causal,
+                  const std::vector<int64_t> window_size, const double softcap,
+                  const bool deterministic, const bool return_attn_probs,
+                  Tensor out, std::optional<Tensor> softmax_lse,
+                  std::optional<Tensor> s_dmask) const override;
+
+ private:
+  void Call(const Tensor q, const Tensor k, const Tensor v,
+            const Tensor cu_seqlens_q, const Tensor cu_seqlens_k,
+            const std::optional<Tensor> alibi_slopes,
+            const std::optional<Tensor> block_table, const int64_t max_seqlen_q,
+            const int64_t max_seqlen_k, const double dropout_p,
+            const std::optional<double> softmax_scale, const bool causal,
+            const std::vector<int64_t> window_size, const double softcap,
+            const bool deterministic, const bool return_attn_probs, Tensor out,
+            std::optional<Tensor> softmax_lse,
+            std::optional<Tensor> s_dmask) const;
+
+  mutable std::mutex runtime_mutex_;
+  mutable linked::tvm_ffi::moore::MateFmhaRuntime runtime_;
+  mutable std::optional<at::Tensor> paged_seqused_k_;
+  mutable std::optional<at::Tensor> internal_lse_;
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_LINKED_TVM_FFI_MOORE_OPS_FLASH_ATTN_VARLEN_FUNC_MATE_H_
diff --git a/src/linked/tvm_ffi/moore/ops/flash_attn_varlen_func/mate.yaml b/src/linked/tvm_ffi/moore/ops/flash_attn_varlen_func/mate.yaml
new file mode 100644
index 000000000..6e76cce0e
--- /dev/null
+++ b/src/linked/tvm_ffi/moore/ops/flash_attn_varlen_func/mate.yaml
@@ -0,0 +1,3 @@
+library: tvm_ffi
+required_symbols:
+  - TVMFFIEnvSetStream
diff --git a/src/linked/tvm_ffi/moore/ops/flash_attn_with_kvcache/mate.cc b/src/linked/tvm_ffi/moore/ops/flash_attn_with_kvcache/mate.cc
new file mode 100644
index 000000000..e6c0e2f98
--- /dev/null
+++ b/src/linked/tvm_ffi/moore/ops/flash_attn_with_kvcache/mate.cc
@@ -0,0 +1,280 @@
+#include "linked/tvm_ffi/moore/ops/flash_attn_with_kvcache/mate.h"
+
+#include <tvm/ffi/container/array.h>
+
+#include <cmath>
+#include <stdexcept>
+
+#include "torch/moore/c10.h"
+#include "torch/tensor_.h"
+
+namespace infini::ops {
+namespace {
+
+namespace mate = linked::tvm_ffi::moore;
+using mate::OptionalTensorView;
+
+py::object PythonTensor(const std::optional<at::Tensor>& tensor) {
+  return tensor.has_value() ? py::cast(*tensor) : py::none();
+}
+
+void CallCombine(mate::MateFmhaRuntime& runtime,
+                 mate::detail::DlPackTensor& out,
+                 mate::detail::DlPackTensor& lse,
+                 const tvm::ffi::Array<tvm::ffi::Tensor>& accumulators,
+                 int max_seqlen_q, int num_splits) {
+  auto* combine = runtime.Combine();
+  if (num_splits > 1 && combine == nullptr) {
+    throw std::runtime_error(
+        "MATE selected split-KV execution without a combine kernel");
+  }
+  if (num_splits <= 1) return;
+
+  tvm::ffi::Function::InvokeExternC(
+      nullptr, combine->Entry(), OptionalTensorView{}, OptionalTensorView{},
+      tvm::ffi::Optional<int>{max_seqlen_q}, tvm::ffi::TensorView(out.Get()),
+      tvm::ffi::TensorView(lse.Get()), tvm::ffi::TensorView(accumulators[0]),
+      tvm::ffi::TensorView(accumulators[1]), OptionalTensorView{},
+      int{num_splits});
+}
+
+}  // namespace
+
+void Operator<FlashAttnWithKvcache, Device::Type::kMoore, 16>::operator()(
+    const Tensor q, Tensor k_cache, Tensor v_cache,
+    const std::optional<Tensor> k, const std::optional<Tensor> v,
+    const std::optional<Tensor> rotary_cos,
+    const std::optional<Tensor> rotary_sin, const int64_t cache_seqlens,
+    const std::optional<Tensor> cache_batch_idx,
+    const std::optional<Tensor> cache_leftpad,
+    const std::optional<Tensor> block_table,
+    const std::optional<Tensor> alibi_slopes,
+    const std::optional<double> softmax_scale, const bool causal,
+    const std::vector<int64_t> window_size, const double softcap,
+    const bool rotary_interleaved, const int64_t num_splits,
+    const bool return_softmax_lse, Tensor out,
+    std::optional<Tensor> softmax_lse) const {
+  Run(q, k_cache, v_cache, k, v, rotary_cos, rotary_sin, std::nullopt,
+      cache_seqlens, cache_batch_idx, cache_leftpad, block_table, alibi_slopes,
+      softmax_scale, causal, window_size, softcap, rotary_interleaved,
+      num_splits, return_softmax_lse, out, softmax_lse);
+}
+
+void Operator<FlashAttnWithKvcache, Device::Type::kMoore, 16>::operator()(
+    const Tensor q, Tensor k_cache, Tensor v_cache,
+    const std::optional<Tensor> k, const std::optional<Tensor> v,
+    const std::optional<Tensor> rotary_cos,
+    const std::optional<Tensor> rotary_sin,
+    const std::optional<Tensor> cache_seqlens,
+    const std::optional<Tensor> cache_batch_idx,
+    const std::optional<Tensor> cache_leftpad,
+    const std::optional<Tensor> block_table,
+    const std::optional<Tensor> alibi_slopes,
+    const std::optional<double> softmax_scale, const bool causal,
+    const std::vector<int64_t> window_size, const double softcap,
+    const bool rotary_interleaved, const int64_t num_splits,
+    const bool return_softmax_lse, Tensor out,
+    std::optional<Tensor> softmax_lse) const {
+  Run(q, k_cache, v_cache, k, v, rotary_cos, rotary_sin, cache_seqlens,
+      std::nullopt, cache_batch_idx, cache_leftpad, block_table, alibi_slopes,
+      softmax_scale, causal, window_size, softcap, rotary_interleaved,
+      num_splits, return_softmax_lse, out, softmax_lse);
+}
+
+void Operator<FlashAttnWithKvcache, Device::Type::kMoore, 16>::Run(
+    const Tensor q, Tensor k_cache, Tensor v_cache,
+    const std::optional<Tensor> k, const std::optional<Tensor> v,
+    const std::optional<Tensor> rotary_cos,
+    const std::optional<Tensor> rotary_sin,
+    const std::optional<Tensor> cache_seqlens,
+    const std::optional<int64_t> scalar_cache_seqlens,
+    const std::optional<Tensor> cache_batch_idx,
+    const std::optional<Tensor> cache_leftpad,
+    const std::optional<Tensor> block_table,
+    const std::optional<Tensor> alibi_slopes,
+    const std::optional<double> softmax_scale, const bool causal,
+    const std::vector<int64_t> window_size, const double softcap,
+    const bool rotary_interleaved, const int64_t num_splits,
+    const bool return_softmax_lse, Tensor out,
+    std::optional<Tensor> softmax_lse) const {
+  (void)return_softmax_lse;
+  TORCH_CHECK(!alibi_slopes.has_value(), "MATE ALiBi is not supported");
+
+  std::lock_guard lock{runtime_mutex_};
+  const typename C10<Device::Type::kMoore>::StreamGuard stream_guard{
+      C10<Device::Type::kMoore>::GetStreamFromExternal(stream_, device_index_)};
+
+  auto at_q =
+      ToAtenTensor<Device::Type::kMoore>(const_cast<void*>(q.data()), q_shape_,
+                                         q_strides_, q_dtype_, device_index_);
+  auto at_k_cache = ToAtenTensor<Device::Type::kMoore>(
+      k_cache.data(), k_cache_shape_, k_cache_strides_, k_cache_dtype_,
+      device_index_);
+  auto at_v_cache = ToAtenTensor<Device::Type::kMoore>(
+      v_cache.data(), v_cache_shape_, v_cache_strides_, v_cache_dtype_,
+      device_index_);
+  auto at_out = ToAtenTensor<Device::Type::kMoore>(
+      out.data(), out_shape_, out_strides_, out_dtype_, device_index_);
+
+  auto optional_tensor = [&](const std::optional<Tensor>& tensor,
+                             const Tensor::Shape& shape,
+                             const Tensor::Strides& strides, DataType dtype) {
+    std::optional<at::Tensor> result;
+    if (tensor.has_value()) {
+      result.emplace(ToAtenTensor<Device::Type::kMoore>(
+          const_cast<void*>(tensor->data()), shape, strides, dtype,
+          device_index_));
+    }
+    return result;
+  };
+
+  const auto at_k = optional_tensor(k, k_shape_, k_strides_, k_dtype_);
+  const auto at_v = optional_tensor(v, v_shape_, v_strides_, v_dtype_);
+  const auto at_rotary_cos = optional_tensor(
+      rotary_cos, rotary_cos_shape_, rotary_cos_strides_, rotary_cos_dtype_);
+  const auto at_rotary_sin = optional_tensor(
+      rotary_sin, rotary_sin_shape_, rotary_sin_strides_, rotary_sin_dtype_);
+  auto at_cache_seqlens =
+      optional_tensor(cache_seqlens, cache_seqlens_shape_,
+                      cache_seqlens_strides_, cache_seqlens_dtype_);
+  const auto at_cache_batch_idx =
+      optional_tensor(cache_batch_idx, cache_batch_idx_shape_,
+                      cache_batch_idx_strides_, cache_batch_idx_dtype_);
+  const auto at_cache_leftpad =
+      optional_tensor(cache_leftpad, cache_leftpad_shape_,
+                      cache_leftpad_strides_, cache_leftpad_dtype_);
+  const auto at_block_table =
+      optional_tensor(block_table, block_table_shape_, block_table_strides_,
+                      block_table_dtype_);
+
+  if (!at_cache_seqlens.has_value() && scalar_cache_seqlens.has_value()) {
+    scalar_cache_seqlens_.emplace(
+        at::full({static_cast<int64_t>(batch_size_)},
+                 static_cast<int64_t>(*scalar_cache_seqlens),
+                 at_out.options().dtype(at::kInt)));
+    at_cache_seqlens = scalar_cache_seqlens_;
+  }
+  if (!softmax_lse.has_value() && !internal_lse_.has_value()) {
+    internal_lse_.emplace(at::empty(
+        {static_cast<int64_t>(q_shape_[0]), static_cast<int64_t>(q_shape_[2]),
+         static_cast<int64_t>(q_shape_[1])},
+        at_out.options().dtype(at::kFloat)));
+  }
+  const auto& at_lse = softmax_lse.has_value() ? *softmax_lse : *internal_lse_;
+  const auto max_seqlen_q = static_cast<int>(q_shape_[1]);
+  const auto pack_gqa = q_shape_[2] != k_cache_shape_[2];
+
+  if (runtime_.Forward().Entry() == nullptr) {
+    py::gil_scoped_acquire gil;
+    try {
+      mate::detail::ModuleRecorder recorder;
+      recorder.Forward().attr("_fmha_fwd")(
+          py::arg("q") = at_q, py::arg("k") = at_k_cache,
+          py::arg("v") = at_v_cache, py::arg("k_new") = PythonTensor(at_k),
+          py::arg("v_new") = PythonTensor(at_v), py::arg("q_v") = py::none(),
+          py::arg("cu_seqlens_q") = py::none(),
+          py::arg("cu_seqlens_k") = py::none(),
+          py::arg("cu_seqlens_k_new") = py::none(),
+          py::arg("seqused_q") = py::none(),
+          py::arg("seqused_k") = PythonTensor(at_cache_seqlens),
+          py::arg("max_seqlen_q") = max_seqlen_q,
+          py::arg("max_seqlen_k") = py::none(),
+          py::arg("page_table") = PythonTensor(at_block_table),
+          py::arg("kv_batch_idx") = PythonTensor(at_cache_batch_idx),
+          py::arg("leftpad_k") = PythonTensor(at_cache_leftpad),
+          py::arg("rotary_cos") = PythonTensor(at_rotary_cos),
+          py::arg("rotary_sin") = PythonTensor(at_rotary_sin),
+          py::arg("seqlens_rotary") = py::none(),
+          py::arg("q_descale") = py::none(), py::arg("k_descale") = py::none(),
+          py::arg("v_descale") = py::none(),
+          py::arg("softmax_scale") = static_cast<float>(softmax_scale.value_or(
+              1.0 / std::sqrt(static_cast<double>(head_size_)))),
+          py::arg("is_causal") = causal,
+          py::arg("window_size_left") = window_size[0],
+          py::arg("window_size_right") = window_size[1],
+          py::arg("attention_chunk") = 0,
+          py::arg("learnable_sink") = py::none(),
+          py::arg("softcap") = static_cast<float>(softcap),
+          py::arg("is_rotary_interleaved") = rotary_interleaved,
+          py::arg("scheduler_metadata") = py::none(),
+          py::arg("num_splits") = num_splits, py::arg("pack_gqa") = pack_gqa,
+          py::arg("mp_margin") = 0, py::arg("return_lse") = true,
+          py::arg("lse") = at_lse, py::arg("out") = at_out,
+          py::arg("cp_world_size") = 1, py::arg("cp_rank") = 0,
+          py::arg("cp_tot_seqused_k") = py::none(), py::arg("only_qv") = false);
+      runtime_.Load(recorder.ForwardName(), recorder.CombineName());
+    } catch (const py::error_already_set& error) {
+      TORCH_CHECK(false, "MATE flash_attn_with_kvcache bootstrap failed: ",
+                  error.what());
+    }
+    return;
+  }
+
+  mate::detail::DlPackTensor dl_q{at_q};
+  mate::detail::DlPackTensor dl_k_cache{at_k_cache};
+  mate::detail::DlPackTensor dl_v_cache{at_v_cache};
+  mate::detail::DlPackTensor dl_out{at_out};
+  mate::detail::DlPackTensor dl_lse{at_lse};
+  std::optional<mate::detail::DlPackTensor> dl_k;
+  std::optional<mate::detail::DlPackTensor> dl_v;
+  std::optional<mate::detail::DlPackTensor> dl_rotary_cos;
+  std::optional<mate::detail::DlPackTensor> dl_rotary_sin;
+  std::optional<mate::detail::DlPackTensor> dl_cache_seqlens;
+  std::optional<mate::detail::DlPackTensor> dl_cache_batch_idx;
+  std::optional<mate::detail::DlPackTensor> dl_cache_leftpad;
+  std::optional<mate::detail::DlPackTensor> dl_block_table;
+  if (at_k.has_value()) dl_k.emplace(*at_k);
+  if (at_v.has_value()) dl_v.emplace(*at_v);
+  if (at_rotary_cos.has_value()) dl_rotary_cos.emplace(*at_rotary_cos);
+  if (at_rotary_sin.has_value()) dl_rotary_sin.emplace(*at_rotary_sin);
+  if (at_cache_seqlens.has_value()) {
+    dl_cache_seqlens.emplace(*at_cache_seqlens);
+  }
+  if (at_cache_batch_idx.has_value()) {
+    dl_cache_batch_idx.emplace(*at_cache_batch_idx);
+  }
+  if (at_cache_leftpad.has_value()) {
+    dl_cache_leftpad.emplace(*at_cache_leftpad);
+  }
+  if (at_block_table.has_value()) dl_block_table.emplace(*at_block_table);
+
+  const mate::detail::TvmStreamGuard tvm_stream_guard{
+      dl_q.Get()->device,
+      C10<Device::Type::kMoore>::GetStreamFromExternal(stream_, device_index_)};
+
+  auto optional_view =
+      [](const std::optional<mate::detail::DlPackTensor>& tensor) {
+        return tensor.has_value()
+                   ? OptionalTensorView{tvm::ffi::TensorView(tensor->Get())}
+                   : OptionalTensorView{};
+      };
+
+  auto result = tvm::ffi::Function::InvokeExternC(
+      nullptr, runtime_.Forward().Entry(), tvm::ffi::TensorView(dl_q.Get()),
+      tvm::ffi::TensorView(dl_k_cache.Get()),
+      tvm::ffi::TensorView(dl_v_cache.Get()), optional_view(dl_k),
+      optional_view(dl_v), OptionalTensorView{}, OptionalTensorView{},
+      OptionalTensorView{}, OptionalTensorView{},
+      optional_view(dl_cache_seqlens), tvm::ffi::Optional<int>{max_seqlen_q},
+      tvm::ffi::Optional<int>{}, optional_view(dl_block_table),
+      optional_view(dl_cache_batch_idx), optional_view(dl_cache_leftpad),
+      optional_view(dl_rotary_cos), optional_view(dl_rotary_sin),
+      OptionalTensorView{}, OptionalTensorView{}, OptionalTensorView{},
+      OptionalTensorView{},
+      static_cast<float>(softmax_scale.value_or(
+          1.0 / std::sqrt(static_cast<double>(head_size_)))),
+      causal, int{window_size[0]}, int{window_size[1]}, int{0},
+      static_cast<float>(softcap), int{0}, int{num_splits},
+      OptionalTensorView{}, OptionalTensorView{},
+      tvm::ffi::TensorView(dl_out.Get()), tvm::ffi::TensorView(dl_lse.Get()),
+      int{1}, int{0}, OptionalTensorView{}, false);
+
+  const auto outputs = result.cast<tvm::ffi::Array<tvm::ffi::Any>>();
+  const auto accumulators =
+      outputs[0].cast<tvm::ffi::Array<tvm::ffi::Tensor>>();
+  const auto actual_num_splits = outputs[1].cast<int>();
+  CallCombine(runtime_, dl_out, dl_lse, accumulators, max_seqlen_q,
+              actual_num_splits);
+}
+
+}  // namespace infini::ops
diff --git a/src/linked/tvm_ffi/moore/ops/flash_attn_with_kvcache/mate.h b/src/linked/tvm_ffi/moore/ops/flash_attn_with_kvcache/mate.h
new file mode 100644
index 000000000..1dbd3f53a
--- /dev/null
+++ b/src/linked/tvm_ffi/moore/ops/flash_attn_with_kvcache/mate.h
@@ -0,0 +1,76 @@
+#ifndef INFINI_OPS_LINKED_TVM_FFI_MOORE_OPS_FLASH_ATTN_WITH_KVCACHE_MATE_H_
+#define INFINI_OPS_LINKED_TVM_FFI_MOORE_OPS_FLASH_ATTN_WITH_KVCACHE_MATE_H_
+
+#include <ATen/core/Tensor.h>
+
+#include <mutex>
+#include <optional>
+
+#include "base/flash_attn_with_kvcache.h"
+#include "linked/tvm_ffi/moore/mate.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<FlashAttnWithKvcache, Device::Type::kMoore, 16>
+    : public FlashAttnWithKvcache {
+ public:
+  using FlashAttnWithKvcache::FlashAttnWithKvcache;
+  using FlashAttnWithKvcache::operator();
+
+  void operator()(const Tensor q, Tensor k_cache, Tensor v_cache,
+                  const std::optional<Tensor> k, const std::optional<Tensor> v,
+                  const std::optional<Tensor> rotary_cos,
+                  const std::optional<Tensor> rotary_sin,
+                  const int64_t cache_seqlens,
+                  const std::optional<Tensor> cache_batch_idx,
+                  const std::optional<Tensor> cache_leftpad,
+                  const std::optional<Tensor> block_table,
+                  const std::optional<Tensor> alibi_slopes,
+                  const std::optional<double> softmax_scale, const bool causal,
+                  const std::vector<int64_t> window_size, const double softcap,
+                  const bool rotary_interleaved, const int64_t num_splits,
+                  const bool return_softmax_lse, Tensor out,
+                  std::optional<Tensor> softmax_lse) const override;
+
+  void operator()(const Tensor q, Tensor k_cache, Tensor v_cache,
+                  const std::optional<Tensor> k, const std::optional<Tensor> v,
+                  const std::optional<Tensor> rotary_cos,
+                  const std::optional<Tensor> rotary_sin,
+                  const std::optional<Tensor> cache_seqlens,
+                  const std::optional<Tensor> cache_batch_idx,
+                  const std::optional<Tensor> cache_leftpad,
+                  const std::optional<Tensor> block_table,
+                  const std::optional<Tensor> alibi_slopes,
+                  const std::optional<double> softmax_scale, const bool causal,
+                  const std::vector<int64_t> window_size, const double softcap,
+                  const bool rotary_interleaved, const int64_t num_splits,
+                  const bool return_softmax_lse, Tensor out,
+                  std::optional<Tensor> softmax_lse) const override;
+
+ private:
+  void Run(const Tensor q, Tensor k_cache, Tensor v_cache,
+           const std::optional<Tensor> k, const std::optional<Tensor> v,
+           const std::optional<Tensor> rotary_cos,
+           const std::optional<Tensor> rotary_sin,
+           const std::optional<Tensor> cache_seqlens,
+           const std::optional<int64_t> scalar_cache_seqlens,
+           const std::optional<Tensor> cache_batch_idx,
+           const std::optional<Tensor> cache_leftpad,
+           const std::optional<Tensor> block_table,
+           const std::optional<Tensor> alibi_slopes,
+           const std::optional<double> softmax_scale, const bool causal,
+           const std::vector<int64_t> window_size, const double softcap,
+           const bool rotary_interleaved, const int64_t num_splits,
+           const bool return_softmax_lse, Tensor out,
+           std::optional<Tensor> softmax_lse) const;
+
+  mutable std::mutex runtime_mutex_;
+  mutable linked::tvm_ffi::moore::MateFmhaRuntime runtime_;
+  mutable std::optional<at::Tensor> scalar_cache_seqlens_;
+  mutable std::optional<at::Tensor> internal_lse_;
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_LINKED_TVM_FFI_MOORE_OPS_FLASH_ATTN_WITH_KVCACHE_MATE_H_
diff --git a/src/linked/tvm_ffi/moore/ops/flash_attn_with_kvcache/mate.yaml b/src/linked/tvm_ffi/moore/ops/flash_attn_with_kvcache/mate.yaml
new file mode 100644
index 000000000..6e76cce0e
--- /dev/null
+++ b/src/linked/tvm_ffi/moore/ops/flash_attn_with_kvcache/mate.yaml
@@ -0,0 +1,3 @@
+library: tvm_ffi
+required_symbols:
+  - TVMFFIEnvSetStream
diff --git a/src/linked/tvm_ffi/moore/tvm_ffi.yaml b/src/linked/tvm_ffi/moore/tvm_ffi.yaml
new file mode 100644
index 000000000..a903ff45a
--- /dev/null
+++ b/src/linked/tvm_ffi/moore/tvm_ffi.yaml
@@ -0,0 +1,4 @@
+python_distribution_package: apache-tvm-ffi
+python_distribution_version: ">=0.1.9,<0.2"
+library_glob: tvm_ffi/lib/libtvm_ffi.so
+include_glob: tvm_ffi/include
diff --git a/tests/test_flash_attn_varlen_func.py b/tests/test_flash_attn_varlen_func.py
index d8d15d7d0..eaff597bc 100644
--- a/tests/test_flash_attn_varlen_func.py
+++ b/tests/test_flash_attn_varlen_func.py
@@ -24,6 +24,7 @@
         ((5, 2), (3, 6), 4, 2, True, (-1, -1), 0.125, False, False),
         ((4, 3), (6, 2), 4, 2, False, (2, 1), None, False, False),
         ((4, 3), (6, 2), 4, 2, True, (2, 1), None, False, False),
+        ((2, 3), (130, 300), 4, 2, True, (-1, -1), None, True, False),
         ((2, 3), (130, 300), 4, 2, True, (-1, -1), None, True, True),
     ),
 )
@@ -56,10 +57,18 @@ def test_flash_attn_varlen_func(
         pytest.skip(
             "FlashAttention requires the NVIDIA, Moore, Cambricon, or Ascend backend"
         )
-    if device == "musa" and window_size != (-1, -1):
+    if device == "musa" and implementation_index == 8 and window_size != (-1, -1):
         pytest.skip("TorchMusa FlashAttention does not support local windows")
-    if device == "musa" and not paged and causal and q_lens != k_lens:
+    if (
+        device == "musa"
+        and implementation_index == 8
+        and not paged
+        and causal
+        and q_lens != k_lens
+    ):
         pytest.skip("TorchMusa causal FlashAttention requires matching Q/K lengths")
+    if device == "musa" and implementation_index == 16 and use_alibi:
+        pytest.skip("Mate does not support ALiBi")
 
     if (
         device == "cuda"
@@ -115,7 +124,11 @@ def test_flash_attn_varlen_func(
         else None
     )
     out = torch.empty_like(q)
-    return_attn_probs = device != "npu" and not paged
+    return_attn_probs = (
+        device != "npu"
+        and not paged
+        and not (device == "musa" and implementation_index == 16)
+    )
     softmax_lse = (
         torch.empty(
             (q.size(1), q.size(0)),
@@ -220,6 +233,9 @@ def test_flash_attn_varlen_func_non_default_stream(device, implementation_index)
     elif device == "npu":
         accelerator = torch.npu
         stream_attribute = "npu_stream"
+    elif device == "musa" and implementation_index == 16:
+        accelerator = torch.musa
+        stream_attribute = "musa_stream"
     else:
         pytest.skip("stream coverage requires an accelerator backend")
     if device == "cuda" and implementation_index == 0:
diff --git a/tests/test_flash_attn_with_kvcache.py b/tests/test_flash_attn_with_kvcache.py
index e62902f87..86a51f039 100644
--- a/tests/test_flash_attn_with_kvcache.py
+++ b/tests/test_flash_attn_with_kvcache.py
@@ -199,10 +199,14 @@ def test_flash_attn_with_kvcache_dense(
     rtol,
     atol,
 ):
-    if device not in ("cuda", "mlu"):
-        pytest.skip("FlashAttention FA2 requires the NVIDIA or Cambricon backend")
+    if device not in ("cuda", "musa", "mlu"):
+        pytest.skip(
+            "FlashAttention FA2 requires the NVIDIA, Moore, or Cambricon backend"
+        )
     if device == "cuda" and implementation_index == 0:
         pytest.skip("Iluvatar native provider supports paged decode only")
+    if device == "musa" and implementation_index != 16:
+        pytest.skip("dense attention requires the Mate linked provider")
 
     batch_size, cache_size = 2, 16
     num_heads, num_kv_heads, head_size = 4, 2, 64
@@ -232,7 +236,7 @@ def test_flash_attn_with_kvcache_dense(
     expected_v_cache = v_cache.clone()
     actual_k_cache = k_cache.clone()
     actual_v_cache = v_cache.clone()
-    if device == "mlu":
+    if device in ("musa", "mlu"):
         expected, expected_softmax_lse = _reference_flash_attn_with_kvcache(
             q,
             expected_k_cache,
@@ -304,8 +308,10 @@ def test_flash_attn_with_kvcache_dense(
 
 @pytest.mark.smoke
 def test_flash_attn_with_kvcache_paged(device, implementation_index):
-    if device not in ("cuda", "mlu"):
-        pytest.skip("FlashAttention FA2 requires the NVIDIA or Cambricon backend")
+    if device not in ("cuda", "musa", "mlu"):
+        pytest.skip(
+            "FlashAttention FA2 requires the NVIDIA, Moore, or Cambricon backend"
+        )
 
     batch_size, page_size = 2, 256
     num_heads, num_kv_heads, head_size = 4, 2, 64
@@ -322,7 +328,7 @@ def test_flash_attn_with_kvcache_paged(device, implementation_index):
     v_cache = torch.randn_like(k_cache)
     cache_seqlens = torch.tensor((130, 300), dtype=torch.int32, device=device)
     block_table = torch.tensor(((0, 1), (2, 3)), dtype=torch.int32, device=device)
-    if device == "mlu" or (device == "cuda" and implementation_index == 0):
+    if device in ("musa", "mlu") or (device == "cuda" and implementation_index == 0):
         expected, _ = _reference_flash_attn_with_kvcache(
             q,
             k_cache,
@@ -466,6 +472,9 @@ def test_flash_attn_with_kvcache_non_default_stream(device, implementation_index
     elif device == "mlu":
         accelerator = torch.mlu
         stream_attribute = "mlu_stream"
+    elif device == "musa" and implementation_index == 16:
+        accelerator = torch.musa
+        stream_attribute = "musa_stream"
     else:
         pytest.skip("stream coverage requires an accelerator backend")
     if device == "cuda" and implementation_index == 0:
@@ -474,7 +483,7 @@ def test_flash_attn_with_kvcache_non_default_stream(device, implementation_index
     q = torch.randn((2, 1, 4, 64), dtype=torch.float16, device=device)
     k_cache = torch.randn((2, 8, 2, 64), dtype=torch.float16, device=device)
     v_cache = torch.randn_like(k_cache)
-    if device == "mlu":
+    if device in ("musa", "mlu"):
         expected, _ = _reference_flash_attn_with_kvcache(q, k_cache, v_cache)
     else:
         expected = _get_flash_attn().flash_attn_with_kvcache(q, k_cache, v_cache)