issue/1415 - feat: support multi_head_attention with moore mate

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)
 """
@@ -38,18 +39,20 @@ def _check_mate_available():
     global _MATE_LOADED
     global flash_attn_with_kvcache
     global get_scheduler_metadata
+    global flash_attn_varlen_func
 
     if not _MATE_VERSION_SUPPORTED:
         installed = _MATE_VERSION or "not installed"
         raise RuntimeError(
             f"mate {_REQUIRED_MATE_VERSION} is required, but found {installed}. "
-            "Please build and install MooreThreads/mate v0.1.3 first."
+            "Please build and install MooreThreads mate v0.1.3 first."
         )
 
     if _MATE_LOADED:
         return
 
     try:
+        from flash_attn import flash_attn_varlen_func as _flash_attn_varlen_func
         from flash_attn import (
             flash_attn_with_kvcache as _flash_attn_with_kvcache,
         )
@@ -61,9 +64,47 @@ def _check_mate_available():
 
     flash_attn_with_kvcache = _flash_attn_with_kvcache
     get_scheduler_metadata = _get_scheduler_metadata
+    flash_attn_varlen_func = _flash_attn_varlen_func
     _MATE_LOADED = True
 
 
+# =============================================================================
+# 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
 # =============================================================================

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 <pybind11/embed.h>
+#include <pybind11/pybind11.h>
+#include <stdexcept>
+#include <torch/csrc/utils/pybind.h>
+
+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<graph::GraphTensor> alibi_slopes;
+    float scale;
+    bool is_causal;
+};
+
+void *plan(Tensor out,
+           const Tensor &q,
+           const Tensor &k,
+           const Tensor &v,
+           std::optional<Tensor> 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>(graph::GraphTensor(*alibi_slopes)) : std::nullopt,
+        scale,
+        is_causal};
+}
+
+void run(void *planned_meta) {
+    auto *p = reinterpret_cast<PlannedMeta *>(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<at::Tensor>());
+    } 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<PlannedMeta **>(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