diff --git a/python/infinicore/ops/moore_mate_flash_attn.py b/python/infinicore/ops/moore_mate_flash_attn.py index e50d0f6a8..a7f875ba1 100644 --- a/python/infinicore/ops/moore_mate_flash_attn.py +++ b/python/infinicore/ops/moore_mate_flash_attn.py @@ -6,6 +6,7 @@ - mate (repo : https://github.com/MooreThreads/mate) Provides three entry points: + - moore_mate_flash_attn_dense: dense attention with layout (batch, seq, heads, dim) - moore_mate_flash_attn_decode: decode with layout (num_blocks, block_size, num_kv_heads, head_size) - moore_mate_flash_attn_prefill: variable-length prefill (used by mha_varlen) """ @@ -13,7 +14,11 @@ import torch try: - from flash_attn import flash_attn_with_kvcache, get_scheduler_metadata + from flash_attn import ( + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_scheduler_metadata, + ) _MATE_AVAILABLE = True except ImportError: @@ -34,6 +39,43 @@ def _check_mate_available(): ) +# ============================================================================= +# Dense kernels +# ============================================================================= + + +@torch.inference_mode() +def moore_mate_flash_attn_dense( + q: torch.Tensor, # (batch, seqlen_q, num_heads, head_size) + k: torch.Tensor, # (batch, seqlen_k, num_kv_heads, head_size) + v: torch.Tensor, # (batch, seqlen_k, num_kv_heads, head_size) + scale: float, + causal: bool, +) -> torch.Tensor: + """Dense MHA/GQA entry point backed by mate flash_attn_varlen_func.""" + _check_mate_available() + + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + + num_heads = q.shape[2] + num_kv_heads = k.shape[2] + if num_heads % num_kv_heads != 0: + raise RuntimeError("query head count must be divisible by key/value head count") + + return flash_attn_varlen_func( + q=q, + k=k, + v=v, + softmax_scale=scale, + causal=causal, + num_splits=-1, + pack_gqa=num_heads != num_kv_heads, + return_softmax_lse=False, + ) + + # ============================================================================= # Decode kernels # ============================================================================= diff --git a/src/infinicore/ops/multi_head_attention/moore/mha_flashattn_moore.cc b/src/infinicore/ops/multi_head_attention/moore/mha_flashattn_moore.cc new file mode 100644 index 000000000..10b336a8f --- /dev/null +++ b/src/infinicore/ops/multi_head_attention/moore/mha_flashattn_moore.cc @@ -0,0 +1,105 @@ +#if defined(ENABLE_MOORE_MATE_FLASH_ATTN) + +#include "infinicore/ops/mha.hpp" + +#include "infinicore/adaptor/aten_adaptor.hpp" + +#include +#include +#include +#include + +namespace infinicore::op::mha_impl::flashattn_moore { + +namespace py = pybind11; + +namespace { +class LocalMUSAStreamGuard { +public: + explicit LocalMUSAStreamGuard(const c10::musa::MUSAStream &s) + : prev_(c10::musa::getCurrentMUSAStream(s.device_index())) { + c10::musa::setCurrentMUSAStream(s); + } + ~LocalMUSAStreamGuard() { + c10::musa::setCurrentMUSAStream(prev_); + } + LocalMUSAStreamGuard(const LocalMUSAStreamGuard &) = delete; + LocalMUSAStreamGuard &operator=(const LocalMUSAStreamGuard &) = delete; + +private: + c10::musa::MUSAStream prev_; +}; +} // namespace + +struct PlannedMeta { + graph::GraphTensor out, q, k, v; + std::optional alibi_slopes; + float scale; + bool is_causal; +}; + +void *plan(Tensor out, + const Tensor &q, + const Tensor &k, + const Tensor &v, + std::optional alibi_slopes, + float scale, + bool is_causal) { + return new PlannedMeta{ + graph::GraphTensor(out), + graph::GraphTensor(q), + graph::GraphTensor(k), + graph::GraphTensor(v), + alibi_slopes ? std::optional(graph::GraphTensor(*alibi_slopes)) : std::nullopt, + scale, + is_causal}; +} + +void run(void *planned_meta) { + auto *p = reinterpret_cast(planned_meta); + if (p->alibi_slopes.has_value()) { + throw std::runtime_error( + "[mha/moore] ALiBi is not supported by mate v0.1.3"); + } + + LocalMUSAStreamGuard guard(infinicore::adaptor::get_musa_stream()); + + auto out = infinicore::adaptor::to_aten_tensor(p->out); + auto q = infinicore::adaptor::to_aten_tensor(p->q); + auto k = infinicore::adaptor::to_aten_tensor(p->k); + auto v = infinicore::adaptor::to_aten_tensor(p->v); + + try { + py::gil_scoped_acquire gil; + py::module_ wrapper = py::module_::import( + "infinicore.ops.moore_mate_flash_attn"); + + py::object result = wrapper.attr("moore_mate_flash_attn_dense")( + py::cast(q), + py::cast(k), + py::cast(v), + p->scale, + p->is_causal); + + out.copy_(result.cast()); + } catch (const py::error_already_set &e) { + throw std::runtime_error( + std::string("[mha/moore] Python error: ") + e.what()); + } +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + MultiheadAttention::plan_dispatcher().registerDevice(Device::Type::MOORE, &plan); + MultiheadAttention::run_dispatcher().registerDevice(Device::Type::MOORE, &run); + MultiheadAttention::cleanup_dispatcher().registerDevice(Device::Type::MOORE, &cleanup); + return true; +}(); + +} // namespace infinicore::op::mha_impl::flashattn_moore + +#endif // ENABLE_MOORE_MATE_FLASH_ATTN