From e1c9f4825ea3a42b6534f5986be6ed6d2423a4fe Mon Sep 17 00:00:00 2001 From: zhushuang Date: Wed, 9 Sep 2026 16:18:05 +0800 Subject: [PATCH] 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)