From e1c9f4825ea3a42b6534f5986be6ed6d2423a4fe Mon Sep 17 00:00:00 2001 From: zhushuang Date: Wed, 9 Sep 2026 16:18:05 +0800 Subject: [PATCH 1/6] feat(moore): support flash-attn via MooreThreads/mate v0.2.5 --- src/CMakeLists.txt | 22 ++ src/linked/torch/moore/mate.h | 32 +++ .../moore/ops/flash_attn_varlen_func/mate.cc | 153 +++++++++++++ .../moore/ops/flash_attn_varlen_func/mate.h | 35 +++ .../moore/ops/flash_attn_with_kvcache/mate.cc | 213 ++++++++++++++++++ .../moore/ops/flash_attn_with_kvcache/mate.h | 53 +++++ tests/test_flash_attn_varlen_func.py | 22 +- tests/test_flash_attn_with_kvcache.py | 23 +- 8 files changed, 543 insertions(+), 10 deletions(-) create mode 100644 src/linked/torch/moore/mate.h create mode 100644 src/linked/torch/moore/ops/flash_attn_varlen_func/mate.cc create mode 100644 src/linked/torch/moore/ops/flash_attn_varlen_func/mate.h create mode 100644 src/linked/torch/moore/ops/flash_attn_with_kvcache/mate.cc create mode 100644 src/linked/torch/moore/ops/flash_attn_with_kvcache/mate.h diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 4e475f3c6..eceb647a8 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -843,6 +843,28 @@ 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) + foreach(_op IN LISTS _moore_mate_linked_ops) + if(NOT _infini_ops_op_allowlist OR _op IN_LIST _infini_ops_op_allowlist) + list(APPEND TORCH_SOURCES + "${CMAKE_CURRENT_SOURCE_DIR}/linked/torch/moore/ops/${_op}/mate.cc") + 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) + set(_infini_ops_linked_uses_torch TRUE) + endif() +endif() if(_infini_ops_linked_uses_tvm_ffi) target_sources(infiniops PRIVATE ${INFINI_OPS_LINKED_TVM_FFI_SOURCES}) diff --git a/src/linked/torch/moore/mate.h b/src/linked/torch/moore/mate.h new file mode 100644 index 000000000..db7f78b9f --- /dev/null +++ b/src/linked/torch/moore/mate.h @@ -0,0 +1,32 @@ +#ifndef INFINI_OPS_LINKED_TORCH_MOORE_MATE_H_ +#define INFINI_OPS_LINKED_TORCH_MOORE_MATE_H_ + +#include +#include + +#include +#include + +namespace infini::ops::linked::torch::moore::detail { + +namespace py = pybind11; + +inline py::module_ ImportMateMhaInterface() { + auto mate = py::module_::import("mate"); + const auto version = py::str(mate.attr("__version__")).cast(); + const auto separator = version.find('+'); + TORCH_CHECK(version.substr(0, separator) == "0.2.5", + "Mate 0.2.5 is required by the Moore FlashAttention provider, " + "but found ", + version); + return py::module_::import("mate.mha_interface"); +} + +inline py::object OptionalTensorToPyObject( + const std::optional& tensor) { + return tensor.has_value() ? py::cast(*tensor) : py::none(); +} + +} // namespace infini::ops::linked::torch::moore::detail + +#endif // INFINI_OPS_LINKED_TORCH_MOORE_MATE_H_ diff --git a/src/linked/torch/moore/ops/flash_attn_varlen_func/mate.cc b/src/linked/torch/moore/ops/flash_attn_varlen_func/mate.cc new file mode 100644 index 000000000..c48383d0b --- /dev/null +++ b/src/linked/torch/moore/ops/flash_attn_varlen_func/mate.cc @@ -0,0 +1,153 @@ +#include "linked/torch/moore/ops/flash_attn_varlen_func/mate.h" + +#include +#include +#include + +#include "linked/torch/moore/mate.h" +#include "linked/torch/ops/flash_attn_varlen_func.h" +#include "torch/moore/c10.h" +namespace infini::ops::linked::torch::moore { +namespace { + +std::vector MateFlashAttnVarlenKernel( + at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor cu_seqlens_q, + at::Tensor cu_seqlens_k, std::optional block_table, + int64_t max_seqlen_q, int64_t max_seqlen_k, double softmax_scale, + bool causal, int64_t window_size_left, int64_t window_size_right, + double softcap) { + namespace py = pybind11; + py::gil_scoped_acquire gil; + + try { + auto mha_interface = detail::ImportMateMhaInterface(); + py::object mate_cu_seqlens_k = py::cast(cu_seqlens_k); + py::object seqused_k = py::none(); + if (block_table.has_value()) { + seqused_k = + py::cast(cu_seqlens_k.slice(0, 1) - cu_seqlens_k.slice(0, 0, -1)); + mate_cu_seqlens_k = py::none(); + } + + auto result = mha_interface.attr("flash_attn_varlen_func")( + py::arg("q") = q, py::arg("k") = k, py::arg("v") = v, + py::arg("cu_seqlens_q") = cu_seqlens_q, + py::arg("cu_seqlens_k") = mate_cu_seqlens_k, + py::arg("seqused_k") = seqused_k, + py::arg("max_seqlen_q") = max_seqlen_q, + py::arg("max_seqlen_k") = max_seqlen_k, + py::arg("page_table") = detail::OptionalTensorToPyObject(block_table), + py::arg("softmax_scale") = softmax_scale, py::arg("causal") = causal, + py::arg("window_size") = + py::make_tuple(window_size_left, window_size_right), + py::arg("softcap") = softcap, py::arg("num_splits") = 0, + py::arg("pack_gqa") = (q.size(-2) != k.size(-2)), + py::arg("deterministic") = false, py::arg("return_softmax_lse") = true, + py::arg("backend") = "auto"); + + auto outputs = result.cast(); + TORCH_CHECK(outputs.size() >= 2, + "Mate flash_attn_varlen_func returned no softmax LSE"); + return {outputs[0].cast(), outputs[1].cast(), + at::empty({0}, q.options())}; + } catch (const py::error_already_set& error) { + TORCH_CHECK(false, "Mate flash_attn_varlen_func failed: ", error.what()); + } +} + +} // namespace +} // namespace infini::ops::linked::torch::moore + +TORCH_LIBRARY_FRAGMENT(infini_moore_mate, m) { + m.def( + "flash_attn_varlen_func(Tensor q, Tensor k, Tensor v, " + "Tensor cu_seqlens_q, Tensor cu_seqlens_k, Tensor? block_table, " + "int max_seqlen_q, int max_seqlen_k, float softmax_scale, bool causal, " + "int window_size_left, int window_size_right, float softcap) -> " + "Tensor[]"); +} + +TORCH_LIBRARY_IMPL(infini_moore_mate, PrivateUse1, m) { + m.impl( + "flash_attn_varlen_func", + TORCH_FN(infini::ops::linked::torch::moore::MateFlashAttnVarlenKernel)); +} + +namespace infini::ops::linked::torch::moore { + +struct MateFlashAttnVarlen : C10 { + static std::vector Call( + at::Tensor& q, const at::Tensor& k, const at::Tensor& v, + std::optional& out, const at::Tensor& cu_seqlens_q, + const at::Tensor& cu_seqlens_k, std::optional& seqused_k, + std::optional& leftpad_k, + std::optional& block_table, + std::optional& alibi_slopes, int max_seqlen_q, + int max_seqlen_k, float dropout_p, float softmax_scale, bool zero_tensors, + bool causal, int window_size_left, int window_size_right, float softcap, + bool return_softmax, std::optional generator) { + TORCH_CHECK(!out.has_value(), "Mate does not accept an output tensor here"); + TORCH_CHECK(!seqused_k.has_value(), "Mate seqused_k is not supported yet"); + TORCH_CHECK(!leftpad_k.has_value(), "Mate leftpad_k is not supported yet"); + TORCH_CHECK(!alibi_slopes.has_value(), "Mate ALiBi is not supported"); + TORCH_CHECK(dropout_p == 0.0f, "Mate attention dropout is not supported"); + TORCH_CHECK(!generator.has_value(), + "Mate attention generators are not supported"); + (void)zero_tensors; + (void)return_softmax; + + static const auto op = c10::Dispatcher::singleton().findSchemaOrThrow( + "infini_moore_mate::flash_attn_varlen_func", ""); + c10::Stack stack; + stack.emplace_back(q); + stack.emplace_back(k); + stack.emplace_back(v); + stack.emplace_back(cu_seqlens_q); + stack.emplace_back(cu_seqlens_k); + stack.emplace_back(block_table.has_value() ? c10::IValue(*block_table) + : c10::IValue()); + stack.emplace_back(static_cast(max_seqlen_q)); + stack.emplace_back(static_cast(max_seqlen_k)); + stack.emplace_back(static_cast(softmax_scale)); + stack.emplace_back(causal); + stack.emplace_back(static_cast(window_size_left)); + stack.emplace_back(static_cast(window_size_right)); + stack.emplace_back(static_cast(softcap)); + op.callBoxed(&stack); + + TORCH_CHECK(stack.size() == 1, + "Mate flash_attn_varlen_func returned an invalid result"); + return stack.front().toTensorVector(); + } +}; + +} // namespace infini::ops::linked::torch::moore + +namespace infini::ops { + +void Operator::operator()( + const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q, + const Tensor cu_seqlens_k, const std::optional alibi_slopes, + const std::optional block_table, const int64_t max_seqlen_q, + const int64_t max_seqlen_k, const double dropout_p, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool deterministic, const bool return_attn_probs, Tensor out, + std::optional softmax_lse, std::optional s_dmask) const { + using Delegate = linked::torch::TorchFlashAttnVarlenFunc< + linked::torch::moore::MateFlashAttnVarlen>; + if (!delegate_) { + delegate_ = std::make_unique( + 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); + } + delegate_->set_stream(stream_); + (*delegate_)(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); +} + +} // namespace infini::ops diff --git a/src/linked/torch/moore/ops/flash_attn_varlen_func/mate.h b/src/linked/torch/moore/ops/flash_attn_varlen_func/mate.h new file mode 100644 index 000000000..c5f3f6219 --- /dev/null +++ b/src/linked/torch/moore/ops/flash_attn_varlen_func/mate.h @@ -0,0 +1,35 @@ +#ifndef INFINI_OPS_LINKED_TORCH_MOORE_OPS_FLASH_ATTN_VARLEN_FUNC_MATE_H_ +#define INFINI_OPS_LINKED_TORCH_MOORE_OPS_FLASH_ATTN_VARLEN_FUNC_MATE_H_ + +#include + +#include "base/flash_attn_varlen_func.h" + +namespace infini::ops { + +template <> +class Operator + : 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 alibi_slopes, + const std::optional block_table, + const int64_t max_seqlen_q, const int64_t max_seqlen_k, + const double dropout_p, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool deterministic, const bool return_attn_probs, + Tensor out, std::optional softmax_lse, + std::optional s_dmask) const override; + + private: + mutable std::unique_ptr delegate_; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_LINKED_TORCH_MOORE_OPS_FLASH_ATTN_VARLEN_FUNC_MATE_H_ diff --git a/src/linked/torch/moore/ops/flash_attn_with_kvcache/mate.cc b/src/linked/torch/moore/ops/flash_attn_with_kvcache/mate.cc new file mode 100644 index 000000000..6034aeb6a --- /dev/null +++ b/src/linked/torch/moore/ops/flash_attn_with_kvcache/mate.cc @@ -0,0 +1,213 @@ +#include "linked/torch/moore/ops/flash_attn_with_kvcache/mate.h" + +#include +#include +#include + +#include "linked/torch/moore/mate.h" +#include "linked/torch/ops/flash_attn_with_kvcache.h" +#include "torch/moore/c10.h" +namespace infini::ops::linked::torch::moore { +namespace { + +std::vector MateFlashAttnKvcacheKernel( + at::Tensor q, at::Tensor k_cache, at::Tensor v_cache, + std::optional k, std::optional v, + std::optional cache_seqlens, + std::optional rotary_cos, std::optional rotary_sin, + std::optional cache_batch_idx, + std::optional cache_leftpad, + std::optional block_table, double softmax_scale, bool causal, + int64_t window_size_left, int64_t window_size_right, double softcap, + bool rotary_interleaved, int64_t num_splits) { + namespace py = pybind11; + py::gil_scoped_acquire gil; + + try { + auto mha_interface = detail::ImportMateMhaInterface(); + auto result = mha_interface.attr("flash_attn_with_kvcache")( + py::arg("q") = q, py::arg("k_cache") = k_cache, + py::arg("v_cache") = v_cache, + py::arg("k") = detail::OptionalTensorToPyObject(k), + py::arg("v") = detail::OptionalTensorToPyObject(v), + py::arg("rotary_cos") = detail::OptionalTensorToPyObject(rotary_cos), + py::arg("rotary_sin") = detail::OptionalTensorToPyObject(rotary_sin), + py::arg("cache_seqlens") = + detail::OptionalTensorToPyObject(cache_seqlens), + py::arg("cache_batch_idx") = + detail::OptionalTensorToPyObject(cache_batch_idx), + py::arg("cache_leftpad") = + detail::OptionalTensorToPyObject(cache_leftpad), + py::arg("page_table") = detail::OptionalTensorToPyObject(block_table), + py::arg("softmax_scale") = softmax_scale, py::arg("causal") = causal, + py::arg("window_size") = + py::make_tuple(window_size_left, window_size_right), + py::arg("softcap") = softcap, + py::arg("rotary_interleaved") = rotary_interleaved, + py::arg("num_splits") = num_splits, + py::arg("pack_gqa") = (q.size(-2) != k_cache.size(-2)), + py::arg("return_softmax_lse") = true); + + auto outputs = result.cast(); + TORCH_CHECK(outputs.size() >= 2, + "Mate flash_attn_with_kvcache returned no softmax LSE"); + return {outputs[0].cast(), outputs[1].cast()}; + } catch (const py::error_already_set& error) { + TORCH_CHECK(false, "Mate flash_attn_with_kvcache failed: ", error.what()); + } +} + +} // namespace +} // namespace infini::ops::linked::torch::moore + +TORCH_LIBRARY_FRAGMENT(infini_moore_mate, m) { + m.def( + "flash_attn_with_kvcache(Tensor q, Tensor k_cache, Tensor v_cache, " + "Tensor? k, Tensor? v, Tensor? cache_seqlens, Tensor? rotary_cos, " + "Tensor? rotary_sin, Tensor? cache_batch_idx, Tensor? cache_leftpad, " + "Tensor? block_table, float softmax_scale, bool causal, " + "int window_size_left, int window_size_right, float softcap, " + "bool rotary_interleaved, int num_splits) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(infini_moore_mate, PrivateUse1, m) { + m.impl( + "flash_attn_with_kvcache", + TORCH_FN(infini::ops::linked::torch::moore::MateFlashAttnKvcacheKernel)); +} + +namespace infini::ops::linked::torch::moore { + +struct MateFlashAttnKvcache : C10 { + static std::vector Call( + at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& k, std::optional& v, + std::optional& cache_seqlens, + std::optional& rotary_cos, + std::optional& rotary_sin, + std::optional& cache_batch_idx, + std::optional& cache_leftpad, + std::optional& block_table, + std::optional& alibi_slopes, std::optional& out, + float softmax_scale, bool causal, int window_size_left, + int window_size_right, float softcap, bool rotary_interleaved, + int num_splits) { + TORCH_CHECK(!alibi_slopes.has_value(), "Mate ALiBi is not supported"); + TORCH_CHECK(!out.has_value(), "Mate does not accept an output tensor here"); + + auto mutable_optional = [](const std::optional& tensor) { + return tensor.has_value() ? std::optional(*tensor) + : std::nullopt; + }; + + static const auto op = c10::Dispatcher::singleton().findSchemaOrThrow( + "infini_moore_mate::flash_attn_with_kvcache", ""); + c10::Stack stack; + stack.emplace_back(q); + stack.emplace_back(k_cache); + stack.emplace_back(v_cache); + const auto at_k = mutable_optional(k); + const auto at_v = mutable_optional(v); + const auto at_cache_seqlens = mutable_optional(cache_seqlens); + const auto at_rotary_cos = mutable_optional(rotary_cos); + const auto at_rotary_sin = mutable_optional(rotary_sin); + const auto at_cache_batch_idx = mutable_optional(cache_batch_idx); + const auto at_cache_leftpad = mutable_optional(cache_leftpad); + stack.emplace_back(at_k.has_value() ? c10::IValue(*at_k) : c10::IValue()); + stack.emplace_back(at_v.has_value() ? c10::IValue(*at_v) : c10::IValue()); + stack.emplace_back(at_cache_seqlens.has_value() + ? c10::IValue(*at_cache_seqlens) + : c10::IValue()); + stack.emplace_back(at_rotary_cos.has_value() ? c10::IValue(*at_rotary_cos) + : c10::IValue()); + stack.emplace_back(at_rotary_sin.has_value() ? c10::IValue(*at_rotary_sin) + : c10::IValue()); + stack.emplace_back(at_cache_batch_idx.has_value() + ? c10::IValue(*at_cache_batch_idx) + : c10::IValue()); + stack.emplace_back(at_cache_leftpad.has_value() + ? c10::IValue(*at_cache_leftpad) + : c10::IValue()); + stack.emplace_back(block_table.has_value() ? c10::IValue(*block_table) + : c10::IValue()); + stack.emplace_back(static_cast(softmax_scale)); + stack.emplace_back(causal); + stack.emplace_back(static_cast(window_size_left)); + stack.emplace_back(static_cast(window_size_right)); + stack.emplace_back(static_cast(softcap)); + stack.emplace_back(rotary_interleaved); + stack.emplace_back(static_cast(num_splits)); + op.callBoxed(&stack); + + TORCH_CHECK(stack.size() == 1, + "Mate flash_attn_with_kvcache returned an invalid result"); + return stack.front().toTensorVector(); + } +}; + +} // namespace infini::ops::linked::torch::moore + +namespace infini::ops { + +void Operator::operator()( + const Tensor q, Tensor k_cache, Tensor v_cache, + const std::optional k, const std::optional v, + const std::optional rotary_cos, + const std::optional rotary_sin, const int64_t cache_seqlens, + const std::optional cache_batch_idx, + const std::optional cache_leftpad, + const std::optional block_table, + const std::optional alibi_slopes, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool rotary_interleaved, const int64_t num_splits, + const bool return_softmax_lse, Tensor out, + std::optional softmax_lse) const { + using Delegate = linked::torch::TorchFlashAttnWithKvcache< + linked::torch::moore::MateFlashAttnKvcache>; + if (!delegate_) { + delegate_ = std::make_unique( + q, k_cache, v_cache, k, v, rotary_cos, rotary_sin, 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); + } + delegate_->set_stream(stream_); + (*delegate_)(q, k_cache, v_cache, k, v, rotary_cos, rotary_sin, 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::operator()( + const Tensor q, Tensor k_cache, Tensor v_cache, + const std::optional k, const std::optional v, + const std::optional rotary_cos, + const std::optional rotary_sin, + const std::optional cache_seqlens, + const std::optional cache_batch_idx, + const std::optional cache_leftpad, + const std::optional block_table, + const std::optional alibi_slopes, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool rotary_interleaved, const int64_t num_splits, + const bool return_softmax_lse, Tensor out, + std::optional softmax_lse) const { + using Delegate = linked::torch::TorchFlashAttnWithKvcache< + linked::torch::moore::MateFlashAttnKvcache>; + if (!delegate_) { + delegate_ = std::make_unique( + q, k_cache, v_cache, k, v, rotary_cos, rotary_sin, 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); + } + delegate_->set_stream(stream_); + (*delegate_)(q, k_cache, v_cache, k, v, rotary_cos, rotary_sin, 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); +} + +} // namespace infini::ops diff --git a/src/linked/torch/moore/ops/flash_attn_with_kvcache/mate.h b/src/linked/torch/moore/ops/flash_attn_with_kvcache/mate.h new file mode 100644 index 000000000..c7af43975 --- /dev/null +++ b/src/linked/torch/moore/ops/flash_attn_with_kvcache/mate.h @@ -0,0 +1,53 @@ +#ifndef INFINI_OPS_LINKED_TORCH_MOORE_OPS_FLASH_ATTN_WITH_KVCACHE_MATE_H_ +#define INFINI_OPS_LINKED_TORCH_MOORE_OPS_FLASH_ATTN_WITH_KVCACHE_MATE_H_ + +#include + +#include "base/flash_attn_with_kvcache.h" + +namespace infini::ops { + +template <> +class Operator + : public FlashAttnWithKvcache { + public: + using FlashAttnWithKvcache::FlashAttnWithKvcache; + using FlashAttnWithKvcache::operator(); + + void operator()(const Tensor q, Tensor k_cache, Tensor v_cache, + const std::optional k, const std::optional v, + const std::optional rotary_cos, + const std::optional rotary_sin, + const int64_t cache_seqlens, + const std::optional cache_batch_idx, + const std::optional cache_leftpad, + const std::optional block_table, + const std::optional alibi_slopes, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool rotary_interleaved, const int64_t num_splits, + const bool return_softmax_lse, Tensor out, + std::optional softmax_lse) const override; + + void operator()(const Tensor q, Tensor k_cache, Tensor v_cache, + const std::optional k, const std::optional v, + const std::optional rotary_cos, + const std::optional rotary_sin, + const std::optional cache_seqlens, + const std::optional cache_batch_idx, + const std::optional cache_leftpad, + const std::optional block_table, + const std::optional alibi_slopes, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool rotary_interleaved, const int64_t num_splits, + const bool return_softmax_lse, Tensor out, + std::optional softmax_lse) const override; + + private: + mutable std::unique_ptr delegate_; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_LINKED_TORCH_MOORE_OPS_FLASH_ATTN_WITH_KVCACHE_MATE_H_ 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) From ffd58b2359bf34e18b62d1b2bb15d00409ea547a Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Thu, 10 Sep 2026 13:18:55 +0800 Subject: [PATCH 2/6] wip: call MATE flash attention natively --- src/CMakeLists.txt | 11 +- src/linked/torch/moore/mate.h | 32 -- .../moore/ops/flash_attn_varlen_func/mate.cc | 153 ---------- .../moore/ops/flash_attn_varlen_func/mate.h | 35 --- .../moore/ops/flash_attn_with_kvcache/mate.cc | 213 ------------- src/linked/tvm_ffi/moore/mate.h | 231 ++++++++++++++ .../moore/ops/flash_attn_varlen_func/mate.cc | 225 ++++++++++++++ .../moore/ops/flash_attn_varlen_func/mate.h | 53 ++++ .../ops/flash_attn_varlen_func/mate.yaml | 3 + .../moore/ops/flash_attn_with_kvcache/mate.cc | 281 ++++++++++++++++++ .../moore/ops/flash_attn_with_kvcache/mate.h | 33 +- .../ops/flash_attn_with_kvcache/mate.yaml | 3 + src/linked/tvm_ffi/moore/tvm_ffi.yaml | 4 + 13 files changed, 835 insertions(+), 442 deletions(-) delete mode 100644 src/linked/torch/moore/mate.h delete mode 100644 src/linked/torch/moore/ops/flash_attn_varlen_func/mate.cc delete mode 100644 src/linked/torch/moore/ops/flash_attn_varlen_func/mate.h delete mode 100644 src/linked/torch/moore/ops/flash_attn_with_kvcache/mate.cc create mode 100644 src/linked/tvm_ffi/moore/mate.h create mode 100644 src/linked/tvm_ffi/moore/ops/flash_attn_varlen_func/mate.cc create mode 100644 src/linked/tvm_ffi/moore/ops/flash_attn_varlen_func/mate.h create mode 100644 src/linked/tvm_ffi/moore/ops/flash_attn_varlen_func/mate.yaml create mode 100644 src/linked/tvm_ffi/moore/ops/flash_attn_with_kvcache/mate.cc rename src/linked/{torch => tvm_ffi}/moore/ops/flash_attn_with_kvcache/mate.h (60%) create mode 100644 src/linked/tvm_ffi/moore/ops/flash_attn_with_kvcache/mate.yaml create mode 100644 src/linked/tvm_ffi/moore/tvm_ffi.yaml diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index eceb647a8..387c8f052 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -849,9 +849,11 @@ if(WITH_MOORE AND WITH_LINKED) flash_attn_with_kvcache) set(_moore_mate_linked_enabled FALSE) foreach(_op IN LISTS _moore_mate_linked_ops) - if(NOT _infini_ops_op_allowlist OR _op IN_LIST _infini_ops_op_allowlist) - list(APPEND TORCH_SOURCES - "${CMAKE_CURRENT_SOURCE_DIR}/linked/torch/moore/ops/${_op}/mate.cc") + 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) set(_moore_mate_linked_enabled TRUE) endif() endforeach() @@ -861,7 +863,8 @@ if(WITH_MOORE AND WITH_LINKED) 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) + list(APPEND TORCH_LIBRARIES + ${TORCH_PYTHON_LIB} Python::Python ${CMAKE_DL_LIBS}) set(_infini_ops_linked_uses_torch TRUE) endif() endif() diff --git a/src/linked/torch/moore/mate.h b/src/linked/torch/moore/mate.h deleted file mode 100644 index db7f78b9f..000000000 --- a/src/linked/torch/moore/mate.h +++ /dev/null @@ -1,32 +0,0 @@ -#ifndef INFINI_OPS_LINKED_TORCH_MOORE_MATE_H_ -#define INFINI_OPS_LINKED_TORCH_MOORE_MATE_H_ - -#include -#include - -#include -#include - -namespace infini::ops::linked::torch::moore::detail { - -namespace py = pybind11; - -inline py::module_ ImportMateMhaInterface() { - auto mate = py::module_::import("mate"); - const auto version = py::str(mate.attr("__version__")).cast(); - const auto separator = version.find('+'); - TORCH_CHECK(version.substr(0, separator) == "0.2.5", - "Mate 0.2.5 is required by the Moore FlashAttention provider, " - "but found ", - version); - return py::module_::import("mate.mha_interface"); -} - -inline py::object OptionalTensorToPyObject( - const std::optional& tensor) { - return tensor.has_value() ? py::cast(*tensor) : py::none(); -} - -} // namespace infini::ops::linked::torch::moore::detail - -#endif // INFINI_OPS_LINKED_TORCH_MOORE_MATE_H_ diff --git a/src/linked/torch/moore/ops/flash_attn_varlen_func/mate.cc b/src/linked/torch/moore/ops/flash_attn_varlen_func/mate.cc deleted file mode 100644 index c48383d0b..000000000 --- a/src/linked/torch/moore/ops/flash_attn_varlen_func/mate.cc +++ /dev/null @@ -1,153 +0,0 @@ -#include "linked/torch/moore/ops/flash_attn_varlen_func/mate.h" - -#include -#include -#include - -#include "linked/torch/moore/mate.h" -#include "linked/torch/ops/flash_attn_varlen_func.h" -#include "torch/moore/c10.h" -namespace infini::ops::linked::torch::moore { -namespace { - -std::vector MateFlashAttnVarlenKernel( - at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor cu_seqlens_q, - at::Tensor cu_seqlens_k, std::optional block_table, - int64_t max_seqlen_q, int64_t max_seqlen_k, double softmax_scale, - bool causal, int64_t window_size_left, int64_t window_size_right, - double softcap) { - namespace py = pybind11; - py::gil_scoped_acquire gil; - - try { - auto mha_interface = detail::ImportMateMhaInterface(); - py::object mate_cu_seqlens_k = py::cast(cu_seqlens_k); - py::object seqused_k = py::none(); - if (block_table.has_value()) { - seqused_k = - py::cast(cu_seqlens_k.slice(0, 1) - cu_seqlens_k.slice(0, 0, -1)); - mate_cu_seqlens_k = py::none(); - } - - auto result = mha_interface.attr("flash_attn_varlen_func")( - py::arg("q") = q, py::arg("k") = k, py::arg("v") = v, - py::arg("cu_seqlens_q") = cu_seqlens_q, - py::arg("cu_seqlens_k") = mate_cu_seqlens_k, - py::arg("seqused_k") = seqused_k, - py::arg("max_seqlen_q") = max_seqlen_q, - py::arg("max_seqlen_k") = max_seqlen_k, - py::arg("page_table") = detail::OptionalTensorToPyObject(block_table), - py::arg("softmax_scale") = softmax_scale, py::arg("causal") = causal, - py::arg("window_size") = - py::make_tuple(window_size_left, window_size_right), - py::arg("softcap") = softcap, py::arg("num_splits") = 0, - py::arg("pack_gqa") = (q.size(-2) != k.size(-2)), - py::arg("deterministic") = false, py::arg("return_softmax_lse") = true, - py::arg("backend") = "auto"); - - auto outputs = result.cast(); - TORCH_CHECK(outputs.size() >= 2, - "Mate flash_attn_varlen_func returned no softmax LSE"); - return {outputs[0].cast(), outputs[1].cast(), - at::empty({0}, q.options())}; - } catch (const py::error_already_set& error) { - TORCH_CHECK(false, "Mate flash_attn_varlen_func failed: ", error.what()); - } -} - -} // namespace -} // namespace infini::ops::linked::torch::moore - -TORCH_LIBRARY_FRAGMENT(infini_moore_mate, m) { - m.def( - "flash_attn_varlen_func(Tensor q, Tensor k, Tensor v, " - "Tensor cu_seqlens_q, Tensor cu_seqlens_k, Tensor? block_table, " - "int max_seqlen_q, int max_seqlen_k, float softmax_scale, bool causal, " - "int window_size_left, int window_size_right, float softcap) -> " - "Tensor[]"); -} - -TORCH_LIBRARY_IMPL(infini_moore_mate, PrivateUse1, m) { - m.impl( - "flash_attn_varlen_func", - TORCH_FN(infini::ops::linked::torch::moore::MateFlashAttnVarlenKernel)); -} - -namespace infini::ops::linked::torch::moore { - -struct MateFlashAttnVarlen : C10 { - static std::vector Call( - at::Tensor& q, const at::Tensor& k, const at::Tensor& v, - std::optional& out, const at::Tensor& cu_seqlens_q, - const at::Tensor& cu_seqlens_k, std::optional& seqused_k, - std::optional& leftpad_k, - std::optional& block_table, - std::optional& alibi_slopes, int max_seqlen_q, - int max_seqlen_k, float dropout_p, float softmax_scale, bool zero_tensors, - bool causal, int window_size_left, int window_size_right, float softcap, - bool return_softmax, std::optional generator) { - TORCH_CHECK(!out.has_value(), "Mate does not accept an output tensor here"); - TORCH_CHECK(!seqused_k.has_value(), "Mate seqused_k is not supported yet"); - TORCH_CHECK(!leftpad_k.has_value(), "Mate leftpad_k is not supported yet"); - TORCH_CHECK(!alibi_slopes.has_value(), "Mate ALiBi is not supported"); - TORCH_CHECK(dropout_p == 0.0f, "Mate attention dropout is not supported"); - TORCH_CHECK(!generator.has_value(), - "Mate attention generators are not supported"); - (void)zero_tensors; - (void)return_softmax; - - static const auto op = c10::Dispatcher::singleton().findSchemaOrThrow( - "infini_moore_mate::flash_attn_varlen_func", ""); - c10::Stack stack; - stack.emplace_back(q); - stack.emplace_back(k); - stack.emplace_back(v); - stack.emplace_back(cu_seqlens_q); - stack.emplace_back(cu_seqlens_k); - stack.emplace_back(block_table.has_value() ? c10::IValue(*block_table) - : c10::IValue()); - stack.emplace_back(static_cast(max_seqlen_q)); - stack.emplace_back(static_cast(max_seqlen_k)); - stack.emplace_back(static_cast(softmax_scale)); - stack.emplace_back(causal); - stack.emplace_back(static_cast(window_size_left)); - stack.emplace_back(static_cast(window_size_right)); - stack.emplace_back(static_cast(softcap)); - op.callBoxed(&stack); - - TORCH_CHECK(stack.size() == 1, - "Mate flash_attn_varlen_func returned an invalid result"); - return stack.front().toTensorVector(); - } -}; - -} // namespace infini::ops::linked::torch::moore - -namespace infini::ops { - -void Operator::operator()( - const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q, - const Tensor cu_seqlens_k, const std::optional alibi_slopes, - const std::optional block_table, const int64_t max_seqlen_q, - const int64_t max_seqlen_k, const double dropout_p, - const std::optional softmax_scale, const bool causal, - const std::vector window_size, const double softcap, - const bool deterministic, const bool return_attn_probs, Tensor out, - std::optional softmax_lse, std::optional s_dmask) const { - using Delegate = linked::torch::TorchFlashAttnVarlenFunc< - linked::torch::moore::MateFlashAttnVarlen>; - if (!delegate_) { - delegate_ = std::make_unique( - 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); - } - delegate_->set_stream(stream_); - (*delegate_)(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); -} - -} // namespace infini::ops diff --git a/src/linked/torch/moore/ops/flash_attn_varlen_func/mate.h b/src/linked/torch/moore/ops/flash_attn_varlen_func/mate.h deleted file mode 100644 index c5f3f6219..000000000 --- a/src/linked/torch/moore/ops/flash_attn_varlen_func/mate.h +++ /dev/null @@ -1,35 +0,0 @@ -#ifndef INFINI_OPS_LINKED_TORCH_MOORE_OPS_FLASH_ATTN_VARLEN_FUNC_MATE_H_ -#define INFINI_OPS_LINKED_TORCH_MOORE_OPS_FLASH_ATTN_VARLEN_FUNC_MATE_H_ - -#include - -#include "base/flash_attn_varlen_func.h" - -namespace infini::ops { - -template <> -class Operator - : 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 alibi_slopes, - const std::optional block_table, - const int64_t max_seqlen_q, const int64_t max_seqlen_k, - const double dropout_p, - const std::optional softmax_scale, const bool causal, - const std::vector window_size, const double softcap, - const bool deterministic, const bool return_attn_probs, - Tensor out, std::optional softmax_lse, - std::optional s_dmask) const override; - - private: - mutable std::unique_ptr delegate_; -}; - -} // namespace infini::ops - -#endif // INFINI_OPS_LINKED_TORCH_MOORE_OPS_FLASH_ATTN_VARLEN_FUNC_MATE_H_ diff --git a/src/linked/torch/moore/ops/flash_attn_with_kvcache/mate.cc b/src/linked/torch/moore/ops/flash_attn_with_kvcache/mate.cc deleted file mode 100644 index 6034aeb6a..000000000 --- a/src/linked/torch/moore/ops/flash_attn_with_kvcache/mate.cc +++ /dev/null @@ -1,213 +0,0 @@ -#include "linked/torch/moore/ops/flash_attn_with_kvcache/mate.h" - -#include -#include -#include - -#include "linked/torch/moore/mate.h" -#include "linked/torch/ops/flash_attn_with_kvcache.h" -#include "torch/moore/c10.h" -namespace infini::ops::linked::torch::moore { -namespace { - -std::vector MateFlashAttnKvcacheKernel( - at::Tensor q, at::Tensor k_cache, at::Tensor v_cache, - std::optional k, std::optional v, - std::optional cache_seqlens, - std::optional rotary_cos, std::optional rotary_sin, - std::optional cache_batch_idx, - std::optional cache_leftpad, - std::optional block_table, double softmax_scale, bool causal, - int64_t window_size_left, int64_t window_size_right, double softcap, - bool rotary_interleaved, int64_t num_splits) { - namespace py = pybind11; - py::gil_scoped_acquire gil; - - try { - auto mha_interface = detail::ImportMateMhaInterface(); - auto result = mha_interface.attr("flash_attn_with_kvcache")( - py::arg("q") = q, py::arg("k_cache") = k_cache, - py::arg("v_cache") = v_cache, - py::arg("k") = detail::OptionalTensorToPyObject(k), - py::arg("v") = detail::OptionalTensorToPyObject(v), - py::arg("rotary_cos") = detail::OptionalTensorToPyObject(rotary_cos), - py::arg("rotary_sin") = detail::OptionalTensorToPyObject(rotary_sin), - py::arg("cache_seqlens") = - detail::OptionalTensorToPyObject(cache_seqlens), - py::arg("cache_batch_idx") = - detail::OptionalTensorToPyObject(cache_batch_idx), - py::arg("cache_leftpad") = - detail::OptionalTensorToPyObject(cache_leftpad), - py::arg("page_table") = detail::OptionalTensorToPyObject(block_table), - py::arg("softmax_scale") = softmax_scale, py::arg("causal") = causal, - py::arg("window_size") = - py::make_tuple(window_size_left, window_size_right), - py::arg("softcap") = softcap, - py::arg("rotary_interleaved") = rotary_interleaved, - py::arg("num_splits") = num_splits, - py::arg("pack_gqa") = (q.size(-2) != k_cache.size(-2)), - py::arg("return_softmax_lse") = true); - - auto outputs = result.cast(); - TORCH_CHECK(outputs.size() >= 2, - "Mate flash_attn_with_kvcache returned no softmax LSE"); - return {outputs[0].cast(), outputs[1].cast()}; - } catch (const py::error_already_set& error) { - TORCH_CHECK(false, "Mate flash_attn_with_kvcache failed: ", error.what()); - } -} - -} // namespace -} // namespace infini::ops::linked::torch::moore - -TORCH_LIBRARY_FRAGMENT(infini_moore_mate, m) { - m.def( - "flash_attn_with_kvcache(Tensor q, Tensor k_cache, Tensor v_cache, " - "Tensor? k, Tensor? v, Tensor? cache_seqlens, Tensor? rotary_cos, " - "Tensor? rotary_sin, Tensor? cache_batch_idx, Tensor? cache_leftpad, " - "Tensor? block_table, float softmax_scale, bool causal, " - "int window_size_left, int window_size_right, float softcap, " - "bool rotary_interleaved, int num_splits) -> Tensor[]"); -} - -TORCH_LIBRARY_IMPL(infini_moore_mate, PrivateUse1, m) { - m.impl( - "flash_attn_with_kvcache", - TORCH_FN(infini::ops::linked::torch::moore::MateFlashAttnKvcacheKernel)); -} - -namespace infini::ops::linked::torch::moore { - -struct MateFlashAttnKvcache : C10 { - static std::vector Call( - at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, - std::optional& k, std::optional& v, - std::optional& cache_seqlens, - std::optional& rotary_cos, - std::optional& rotary_sin, - std::optional& cache_batch_idx, - std::optional& cache_leftpad, - std::optional& block_table, - std::optional& alibi_slopes, std::optional& out, - float softmax_scale, bool causal, int window_size_left, - int window_size_right, float softcap, bool rotary_interleaved, - int num_splits) { - TORCH_CHECK(!alibi_slopes.has_value(), "Mate ALiBi is not supported"); - TORCH_CHECK(!out.has_value(), "Mate does not accept an output tensor here"); - - auto mutable_optional = [](const std::optional& tensor) { - return tensor.has_value() ? std::optional(*tensor) - : std::nullopt; - }; - - static const auto op = c10::Dispatcher::singleton().findSchemaOrThrow( - "infini_moore_mate::flash_attn_with_kvcache", ""); - c10::Stack stack; - stack.emplace_back(q); - stack.emplace_back(k_cache); - stack.emplace_back(v_cache); - const auto at_k = mutable_optional(k); - const auto at_v = mutable_optional(v); - const auto at_cache_seqlens = mutable_optional(cache_seqlens); - const auto at_rotary_cos = mutable_optional(rotary_cos); - const auto at_rotary_sin = mutable_optional(rotary_sin); - const auto at_cache_batch_idx = mutable_optional(cache_batch_idx); - const auto at_cache_leftpad = mutable_optional(cache_leftpad); - stack.emplace_back(at_k.has_value() ? c10::IValue(*at_k) : c10::IValue()); - stack.emplace_back(at_v.has_value() ? c10::IValue(*at_v) : c10::IValue()); - stack.emplace_back(at_cache_seqlens.has_value() - ? c10::IValue(*at_cache_seqlens) - : c10::IValue()); - stack.emplace_back(at_rotary_cos.has_value() ? c10::IValue(*at_rotary_cos) - : c10::IValue()); - stack.emplace_back(at_rotary_sin.has_value() ? c10::IValue(*at_rotary_sin) - : c10::IValue()); - stack.emplace_back(at_cache_batch_idx.has_value() - ? c10::IValue(*at_cache_batch_idx) - : c10::IValue()); - stack.emplace_back(at_cache_leftpad.has_value() - ? c10::IValue(*at_cache_leftpad) - : c10::IValue()); - stack.emplace_back(block_table.has_value() ? c10::IValue(*block_table) - : c10::IValue()); - stack.emplace_back(static_cast(softmax_scale)); - stack.emplace_back(causal); - stack.emplace_back(static_cast(window_size_left)); - stack.emplace_back(static_cast(window_size_right)); - stack.emplace_back(static_cast(softcap)); - stack.emplace_back(rotary_interleaved); - stack.emplace_back(static_cast(num_splits)); - op.callBoxed(&stack); - - TORCH_CHECK(stack.size() == 1, - "Mate flash_attn_with_kvcache returned an invalid result"); - return stack.front().toTensorVector(); - } -}; - -} // namespace infini::ops::linked::torch::moore - -namespace infini::ops { - -void Operator::operator()( - const Tensor q, Tensor k_cache, Tensor v_cache, - const std::optional k, const std::optional v, - const std::optional rotary_cos, - const std::optional rotary_sin, const int64_t cache_seqlens, - const std::optional cache_batch_idx, - const std::optional cache_leftpad, - const std::optional block_table, - const std::optional alibi_slopes, - const std::optional softmax_scale, const bool causal, - const std::vector window_size, const double softcap, - const bool rotary_interleaved, const int64_t num_splits, - const bool return_softmax_lse, Tensor out, - std::optional softmax_lse) const { - using Delegate = linked::torch::TorchFlashAttnWithKvcache< - linked::torch::moore::MateFlashAttnKvcache>; - if (!delegate_) { - delegate_ = std::make_unique( - q, k_cache, v_cache, k, v, rotary_cos, rotary_sin, 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); - } - delegate_->set_stream(stream_); - (*delegate_)(q, k_cache, v_cache, k, v, rotary_cos, rotary_sin, 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::operator()( - const Tensor q, Tensor k_cache, Tensor v_cache, - const std::optional k, const std::optional v, - const std::optional rotary_cos, - const std::optional rotary_sin, - const std::optional cache_seqlens, - const std::optional cache_batch_idx, - const std::optional cache_leftpad, - const std::optional block_table, - const std::optional alibi_slopes, - const std::optional softmax_scale, const bool causal, - const std::vector window_size, const double softcap, - const bool rotary_interleaved, const int64_t num_splits, - const bool return_softmax_lse, Tensor out, - std::optional softmax_lse) const { - using Delegate = linked::torch::TorchFlashAttnWithKvcache< - linked::torch::moore::MateFlashAttnKvcache>; - if (!delegate_) { - delegate_ = std::make_unique( - q, k_cache, v_cache, k, v, rotary_cos, rotary_sin, 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); - } - delegate_->set_stream(stream_); - (*delegate_)(q, k_cache, v_cache, k, v, rotary_cos, rotary_sin, 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); -} - -} // namespace infini::ops diff --git a/src/linked/tvm_ffi/moore/mate.h b/src/linked/tvm_ffi/moore/mate.h new file mode 100644 index 000000000..e99c116e3 --- /dev/null +++ b/src/linked/tvm_ffi/moore/mate.h @@ -0,0 +1,231 @@ +#ifndef INFINI_OPS_LINKED_TVM_FFI_MOORE_MATE_H_ +#define INFINI_OPS_LINKED_TVM_FFI_MOORE_MATE_H_ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace infini::ops::linked::tvm_ffi::moore { + +namespace py = pybind11; + +using TvmFfiEntry = int(void*, const TVMFFIAny*, int32_t, TVMFFIAny*); +using OptionalTensorView = tvm::ffi::Optional; + +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; + +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( + 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(); + 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 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(); + } + + 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& 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..0627b2d84 --- /dev/null +++ b/src/linked/tvm_ffi/moore/ops/flash_attn_varlen_func/mate.cc @@ -0,0 +1,225 @@ +#include "linked/tvm_ffi/moore/ops/flash_attn_varlen_func/mate.h" + +#include + +#include +#include +#include + +#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& 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& 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{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::operator()( + const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q, + const Tensor cu_seqlens_k, const std::optional alibi_slopes, + const std::optional block_table, const int64_t max_seqlen_q, + const int64_t max_seqlen_k, const double dropout_p, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool deterministic, const bool return_attn_probs, Tensor out, + std::optional softmax_lse, std::optional 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::Call( + const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q, + const Tensor cu_seqlens_k, const std::optional alibi_slopes, + const std::optional block_table, const int64_t max_seqlen_q, + const int64_t max_seqlen_k, const double dropout_p, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool deterministic, const bool return_attn_probs, Tensor out, + std::optional softmax_lse, std::optional 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::StreamGuard stream_guard{ + C10::GetStreamFromExternal(stream_, device_index_)}; + + auto at_q = ToAtenTensor( + const_cast(q.data()), q_shape_, q_strides_, q_dtype_, + device_index_); + auto at_k = ToAtenTensor( + const_cast(k.data()), k_shape_, k_strides_, k_dtype_, + device_index_); + auto at_v = ToAtenTensor( + const_cast(v.data()), v_shape_, v_strides_, v_dtype_, + device_index_); + auto at_cu_seqlens_q = ToAtenTensor( + const_cast(cu_seqlens_q.data()), cu_seqlens_q_shape_, + cu_seqlens_q_strides_, cu_seqlens_q_dtype_, device_index_); + auto at_cu_seqlens_k = ToAtenTensor( + const_cast(cu_seqlens_k.data()), cu_seqlens_k_shape_, + cu_seqlens_k_strides_, cu_seqlens_k_dtype_, device_index_); + auto at_out = ToAtenTensor( + out.data(), out_shape_, out_strides_, out_dtype_, device_index_); + + std::optional at_block_table; + if (block_table.has_value()) { + at_block_table.emplace(ToAtenTensor( + const_cast(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(q_shape_[1]), static_cast(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(softmax_scale.value_or( + 1.0 / std::sqrt(static_cast(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(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 dl_block_table; + std::optional 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::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{static_cast(max_seqlen_q)}, + tvm::ffi::Optional{static_cast(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(softmax_scale.value_or( + 1.0 / std::sqrt(static_cast(q_shape_[2])))), + causal, int{window_size[0]}, int{window_size[1]}, int{0}, + static_cast(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>(); + const auto accumulators = + outputs[0].cast>(); + const auto num_splits = outputs[1].cast(); + CallCombine(runtime_, dl_out, dl_lse, accumulators, dl_cu_q, + static_cast(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..a5e3bcd20 --- /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 + +#include +#include + +#include "base/flash_attn_varlen_func.h" +#include "linked/tvm_ffi/moore/mate.h" + +namespace infini::ops { + +template <> +class Operator + : 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 alibi_slopes, + const std::optional block_table, + const int64_t max_seqlen_q, const int64_t max_seqlen_k, + const double dropout_p, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool deterministic, const bool return_attn_probs, + Tensor out, std::optional softmax_lse, + std::optional 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 alibi_slopes, + const std::optional block_table, const int64_t max_seqlen_q, + const int64_t max_seqlen_k, const double dropout_p, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool deterministic, const bool return_attn_probs, + Tensor out, std::optional softmax_lse, + std::optional s_dmask) const; + + mutable std::mutex runtime_mutex_; + mutable linked::tvm_ffi::moore::MateFmhaRuntime runtime_; + mutable std::optional paged_seqused_k_; + mutable std::optional 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..8fabf2f83 --- /dev/null +++ b/src/linked/tvm_ffi/moore/ops/flash_attn_with_kvcache/mate.cc @@ -0,0 +1,281 @@ +#include "linked/tvm_ffi/moore/ops/flash_attn_with_kvcache/mate.h" + +#include + +#include +#include + +#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& 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& 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{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::operator()( + const Tensor q, Tensor k_cache, Tensor v_cache, const std::optional k, + const std::optional v, const std::optional rotary_cos, + const std::optional rotary_sin, const int64_t cache_seqlens, + const std::optional cache_batch_idx, + const std::optional cache_leftpad, + const std::optional block_table, + const std::optional alibi_slopes, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool rotary_interleaved, const int64_t num_splits, + const bool return_softmax_lse, Tensor out, + std::optional 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::operator()( + const Tensor q, Tensor k_cache, Tensor v_cache, const std::optional k, + const std::optional v, const std::optional rotary_cos, + const std::optional rotary_sin, + const std::optional cache_seqlens, + const std::optional cache_batch_idx, + const std::optional cache_leftpad, + const std::optional block_table, + const std::optional alibi_slopes, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool rotary_interleaved, const int64_t num_splits, + const bool return_softmax_lse, Tensor out, + std::optional 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::Run( + const Tensor q, Tensor k_cache, Tensor v_cache, const std::optional k, + const std::optional v, const std::optional rotary_cos, + const std::optional rotary_sin, + const std::optional cache_seqlens, + const std::optional scalar_cache_seqlens, + const std::optional cache_batch_idx, + const std::optional cache_leftpad, + const std::optional block_table, + const std::optional alibi_slopes, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool rotary_interleaved, const int64_t num_splits, + const bool return_softmax_lse, Tensor out, + std::optional 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::StreamGuard stream_guard{ + C10::GetStreamFromExternal(stream_, device_index_)}; + + auto at_q = ToAtenTensor( + const_cast(q.data()), q_shape_, q_strides_, q_dtype_, + device_index_); + auto at_k_cache = ToAtenTensor( + k_cache.data(), k_cache_shape_, k_cache_strides_, k_cache_dtype_, + device_index_); + auto at_v_cache = ToAtenTensor( + v_cache.data(), v_cache_shape_, v_cache_strides_, v_cache_dtype_, + device_index_); + auto at_out = ToAtenTensor( + out.data(), out_shape_, out_strides_, out_dtype_, device_index_); + + auto optional_tensor = [&](const std::optional& tensor, + const Tensor::Shape& shape, + const Tensor::Strides& strides, DataType dtype) { + std::optional result; + if (tensor.has_value()) { + result.emplace(ToAtenTensor( + const_cast(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(batch_size_)}, + static_cast(*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(q_shape_[0]), static_cast(q_shape_[2]), + static_cast(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(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(softmax_scale.value_or( + 1.0 / std::sqrt(static_cast(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(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 dl_k; + std::optional dl_v; + std::optional dl_rotary_cos; + std::optional dl_rotary_sin; + std::optional dl_cache_seqlens; + std::optional dl_cache_batch_idx; + std::optional dl_cache_leftpad; + std::optional 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::GetStreamFromExternal(stream_, device_index_)}; + + auto optional_view = + [](const std::optional& 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{max_seqlen_q}, tvm::ffi::Optional{}, + 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(softmax_scale.value_or( + 1.0 / std::sqrt(static_cast(head_size_)))), + causal, int{window_size[0]}, int{window_size[1]}, int{0}, + static_cast(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>(); + const auto accumulators = + outputs[0].cast>(); + const auto actual_num_splits = outputs[1].cast(); + CallCombine(runtime_, dl_out, dl_lse, accumulators, max_seqlen_q, + actual_num_splits); +} + +} // namespace infini::ops diff --git a/src/linked/torch/moore/ops/flash_attn_with_kvcache/mate.h b/src/linked/tvm_ffi/moore/ops/flash_attn_with_kvcache/mate.h similarity index 60% rename from src/linked/torch/moore/ops/flash_attn_with_kvcache/mate.h rename to src/linked/tvm_ffi/moore/ops/flash_attn_with_kvcache/mate.h index c7af43975..1dbd3f53a 100644 --- a/src/linked/torch/moore/ops/flash_attn_with_kvcache/mate.h +++ b/src/linked/tvm_ffi/moore/ops/flash_attn_with_kvcache/mate.h @@ -1,9 +1,13 @@ -#ifndef INFINI_OPS_LINKED_TORCH_MOORE_OPS_FLASH_ATTN_WITH_KVCACHE_MATE_H_ -#define INFINI_OPS_LINKED_TORCH_MOORE_OPS_FLASH_ATTN_WITH_KVCACHE_MATE_H_ +#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 +#include + +#include +#include #include "base/flash_attn_with_kvcache.h" +#include "linked/tvm_ffi/moore/mate.h" namespace infini::ops { @@ -45,9 +49,28 @@ class Operator std::optional softmax_lse) const override; private: - mutable std::unique_ptr delegate_; + void Run(const Tensor q, Tensor k_cache, Tensor v_cache, + const std::optional k, const std::optional v, + const std::optional rotary_cos, + const std::optional rotary_sin, + const std::optional cache_seqlens, + const std::optional scalar_cache_seqlens, + const std::optional cache_batch_idx, + const std::optional cache_leftpad, + const std::optional block_table, + const std::optional alibi_slopes, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool rotary_interleaved, const int64_t num_splits, + const bool return_softmax_lse, Tensor out, + std::optional softmax_lse) const; + + mutable std::mutex runtime_mutex_; + mutable linked::tvm_ffi::moore::MateFmhaRuntime runtime_; + mutable std::optional scalar_cache_seqlens_; + mutable std::optional internal_lse_; }; } // namespace infini::ops -#endif // INFINI_OPS_LINKED_TORCH_MOORE_OPS_FLASH_ATTN_WITH_KVCACHE_MATE_H_ +#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 From 62c5b83bc59d3ee3ebd22401c3cf9828c886e3b2 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Thu, 10 Sep 2026 14:20:32 +0800 Subject: [PATCH 3/6] fix: compile Moore MATE glue with system CXX --- src/CMakeLists.txt | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 387c8f052..3da3008fa 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -848,12 +848,14 @@ if(WITH_MOORE AND WITH_LINKED) 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() @@ -866,6 +868,8 @@ if(WITH_MOORE AND WITH_LINKED) 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() From cf2b37c7ce6fb25d89a4955b68f4e63763391387 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Thu, 10 Sep 2026 15:36:35 +0800 Subject: [PATCH 4/6] fix: include linked headers in system CXX build --- src/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 3da3008fa..6d64f6b08 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -1039,7 +1039,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() From 86e2000949e56a238a28e258296032eee591d857 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 15 Sep 2026 13:30:05 +0800 Subject: [PATCH 5/6] style: format Moore MATE provider --- src/linked/tvm_ffi/moore/mate.h | 30 ++++---- .../moore/ops/flash_attn_varlen_func/mate.cc | 45 ++++++------ .../moore/ops/flash_attn_varlen_func/mate.h | 4 +- .../moore/ops/flash_attn_with_kvcache/mate.cc | 69 +++++++++---------- 4 files changed, 70 insertions(+), 78 deletions(-) diff --git a/src/linked/tvm_ffi/moore/mate.h b/src/linked/tvm_ffi/moore/mate.h index e99c116e3..1ad96bab9 100644 --- a/src/linked/tvm_ffi/moore/mate.h +++ b/src/linked/tvm_ffi/moore/mate.h @@ -10,12 +10,12 @@ #include #include -#include #include +#include +#include #include #include #include -#include #include namespace infini::ops::linked::tvm_ffi::moore { @@ -51,8 +51,8 @@ class DlSymbol { entry_ = reinterpret_cast( 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); + throw std::runtime_error("MATE did not export __tvm_ffi_" + + dispatch_name + ": " + error); } } @@ -104,9 +104,8 @@ class DlPackTensor { class TvmStreamGuard { public: TvmStreamGuard(DLDevice device, void* stream) : device_{device} { - const auto status = - TVMFFIEnvSetStream(device.device_type, device.device_id, stream, - &previous_stream_); + 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 " + @@ -137,15 +136,14 @@ class ModuleRecorder { const auto version = py::str(mate.attr("__version__")).cast(); 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); + 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"); + 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(); @@ -179,9 +177,7 @@ class ModuleRecorder { py::module_ Forward() const { return forward_; } - std::string ForwardName() const { - return Name(forward_names_); - } + std::string ForwardName() const { return Name(forward_names_); } std::optional CombineName() const { if (py::len(combine_names_) == 0) return std::nullopt; 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 index 0627b2d84..a611adad8 100644 --- 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 @@ -56,8 +56,8 @@ void Operator::operator()( 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); + window_size, softcap, deterministic, return_attn_probs, out, softmax_lse, + s_dmask); } void Operator::Call( @@ -73,20 +73,21 @@ void Operator::Call( 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"); + TORCH_CHECK(!return_attn_probs, + "MATE attention probabilities are not supported"); const typename C10::StreamGuard stream_guard{ C10::GetStreamFromExternal(stream_, device_index_)}; - auto at_q = ToAtenTensor( - const_cast(q.data()), q_shape_, q_strides_, q_dtype_, - device_index_); - auto at_k = ToAtenTensor( - const_cast(k.data()), k_shape_, k_strides_, k_dtype_, - device_index_); - auto at_v = ToAtenTensor( - const_cast(v.data()), v_shape_, v_strides_, v_dtype_, - device_index_); + auto at_q = + ToAtenTensor(const_cast(q.data()), q_shape_, + q_strides_, q_dtype_, device_index_); + auto at_k = + ToAtenTensor(const_cast(k.data()), k_shape_, + k_strides_, k_dtype_, device_index_); + auto at_v = + ToAtenTensor(const_cast(v.data()), v_shape_, + v_strides_, v_dtype_, device_index_); auto at_cu_seqlens_q = ToAtenTensor( const_cast(cu_seqlens_q.data()), cu_seqlens_q_shape_, cu_seqlens_q_strides_, cu_seqlens_q_dtype_, device_index_); @@ -134,12 +135,10 @@ void Operator::Call( 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("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("q_descale") = py::none(), py::arg("k_descale") = py::none(), py::arg("v_descale") = py::none(), py::arg("softmax_scale") = static_cast(softmax_scale.value_or( 1.0 / std::sqrt(static_cast(q_shape_[2])))), @@ -150,16 +149,14 @@ void Operator::Call( py::arg("learnable_sink") = py::none(), py::arg("softcap") = static_cast(softcap), py::arg("is_rotary_interleaved") = false, - py::arg("scheduler_metadata") = py::none(), - py::arg("num_splits") = 0, + 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])), + (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); + 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: ", @@ -211,8 +208,8 @@ void Operator::Call( causal, int{window_size[0]}, int{window_size[1]}, int{0}, static_cast(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); + tvm::ffi::TensorView(dl_lse.Get()), int{1}, int{0}, OptionalTensorView{}, + false); const auto outputs = result.cast>(); const auto accumulators = 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 index a5e3bcd20..5b1200889 100644 --- 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 @@ -38,8 +38,8 @@ class Operator const int64_t max_seqlen_k, const double dropout_p, const std::optional softmax_scale, const bool causal, const std::vector window_size, const double softcap, - const bool deterministic, const bool return_attn_probs, - Tensor out, std::optional softmax_lse, + const bool deterministic, const bool return_attn_probs, Tensor out, + std::optional softmax_lse, std::optional s_dmask) const; mutable std::mutex runtime_mutex_; 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 index 8fabf2f83..e6c0e2f98 100644 --- 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 @@ -32,9 +32,8 @@ void CallCombine(mate::MateFmhaRuntime& runtime, tvm::ffi::Function::InvokeExternC( nullptr, combine->Entry(), OptionalTensorView{}, OptionalTensorView{}, - tvm::ffi::Optional{max_seqlen_q}, - tvm::ffi::TensorView(out.Get()), tvm::ffi::TensorView(lse.Get()), - tvm::ffi::TensorView(accumulators[0]), + tvm::ffi::Optional{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}); } @@ -42,8 +41,9 @@ void CallCombine(mate::MateFmhaRuntime& runtime, } // namespace void Operator::operator()( - const Tensor q, Tensor k_cache, Tensor v_cache, const std::optional k, - const std::optional v, const std::optional rotary_cos, + const Tensor q, Tensor k_cache, Tensor v_cache, + const std::optional k, const std::optional v, + const std::optional rotary_cos, const std::optional rotary_sin, const int64_t cache_seqlens, const std::optional cache_batch_idx, const std::optional cache_leftpad, @@ -55,14 +55,15 @@ void Operator::operator()( const bool return_softmax_lse, Tensor out, std::optional 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); + 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::operator()( - const Tensor q, Tensor k_cache, Tensor v_cache, const std::optional k, - const std::optional v, const std::optional rotary_cos, + const Tensor q, Tensor k_cache, Tensor v_cache, + const std::optional k, const std::optional v, + const std::optional rotary_cos, const std::optional rotary_sin, const std::optional cache_seqlens, const std::optional cache_batch_idx, @@ -75,14 +76,15 @@ void Operator::operator()( const bool return_softmax_lse, Tensor out, std::optional 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); + 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::Run( - const Tensor q, Tensor k_cache, Tensor v_cache, const std::optional k, - const std::optional v, const std::optional rotary_cos, + const Tensor q, Tensor k_cache, Tensor v_cache, + const std::optional k, const std::optional v, + const std::optional rotary_cos, const std::optional rotary_sin, const std::optional cache_seqlens, const std::optional scalar_cache_seqlens, @@ -102,9 +104,9 @@ void Operator::Run( const typename C10::StreamGuard stream_guard{ C10::GetStreamFromExternal(stream_, device_index_)}; - auto at_q = ToAtenTensor( - const_cast(q.data()), q_shape_, q_strides_, q_dtype_, - device_index_); + auto at_q = + ToAtenTensor(const_cast(q.data()), q_shape_, + q_strides_, q_dtype_, device_index_); auto at_k_cache = ToAtenTensor( k_cache.data(), k_cache_shape_, k_cache_strides_, k_cache_dtype_, device_index_); @@ -142,14 +144,14 @@ void Operator::Run( 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_); + 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(batch_size_)}, - static_cast(*scalar_cache_seqlens), - at_out.options().dtype(at::kInt))); + scalar_cache_seqlens_.emplace( + at::full({static_cast(batch_size_)}, + static_cast(*scalar_cache_seqlens), + at_out.options().dtype(at::kInt))); at_cache_seqlens = scalar_cache_seqlens_; } if (!softmax_lse.has_value() && !internal_lse_.has_value()) { @@ -183,8 +185,7 @@ void Operator::Run( 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("q_descale") = py::none(), py::arg("k_descale") = py::none(), py::arg("v_descale") = py::none(), py::arg("softmax_scale") = static_cast(softmax_scale.value_or( 1.0 / std::sqrt(static_cast(head_size_)))), @@ -196,13 +197,11 @@ void Operator::Run( py::arg("softcap") = static_cast(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("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); + 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: ", @@ -256,12 +255,12 @@ void Operator::Run( 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{max_seqlen_q}, tvm::ffi::Optional{}, - 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{}, + optional_view(dl_cache_seqlens), tvm::ffi::Optional{max_seqlen_q}, + tvm::ffi::Optional{}, 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(softmax_scale.value_or( 1.0 / std::sqrt(static_cast(head_size_)))), causal, int{window_size[0]}, int{window_size[1]}, int{0}, From 8d6e7e6028351275edd53ef4ad6c8a7e64d490fa Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 15 Sep 2026 13:43:19 +0800 Subject: [PATCH 6/6] fix: compile MATE providers separately --- src/CMakeLists.txt | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 6d64f6b08..24ebad24f 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -958,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}") @@ -968,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") @@ -1003,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 $)