issue/1415 - feat: support multi_head_attention with moore mate

python/infinicore/ops/moore_mate_flash_attn.py

@@ -6,14 +6,19 @@
     - 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)
 """
 
 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
 # =============================================================================

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