feat(linked): add MetaX FlashAttention providers

src/linked/torch/metax/flash_attn.yaml

@@ -0,0 +1,2 @@
+python_distribution_package: flash-attn
+library_glob: flash_attn_2_cuda*.so

src/linked/torch/metax/ops/flash_attn_varlen_func/flash_attn.cc

@@ -0,0 +1,72 @@
+#include "linked/torch/metax/ops/flash_attn_varlen_func/flash_attn.h"
+
+#include <ATen/core/Generator.h>
+
+#include "linked/torch/ops/flash_attn_varlen_func.h"
+#include "torch/metax/c10.h"
+
+std::vector<at::Tensor> mha_varlen_fwd(
+    at::Tensor& q, const at::Tensor& k, const at::Tensor& v,
+    std::optional<at::Tensor>& out, const at::Tensor& cu_seqlens_q,
+    const at::Tensor& cu_seqlens_k, std::optional<at::Tensor>& seqused_k,
+    std::optional<const at::Tensor>& leftpad_k,
+    std::optional<at::Tensor>& block_table,
+    std::optional<at::Tensor>& 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<at::Generator> generator,
+    std::optional<at::Tensor>& flash_attn_mars_ext);
+
+namespace infini::ops::linked::torch::metax {
+
+struct FlashAttnVarlen : C10<Device::Type::kMetax> {
+  static std::vector<at::Tensor> Call(
+      at::Tensor& q, const at::Tensor& k, const at::Tensor& v,
+      std::optional<at::Tensor>& out, const at::Tensor& cu_seqlens_q,
+      const at::Tensor& cu_seqlens_k, std::optional<at::Tensor>& seqused_k,
+      std::optional<const at::Tensor>& leftpad_k,
+      std::optional<at::Tensor>& block_table,
+      std::optional<at::Tensor>& 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<at::Generator> generator) {
+    std::optional<at::Tensor> flash_attn_mars_ext;
+    return ::mha_varlen_fwd(q, k, v, out, cu_seqlens_q, cu_seqlens_k, seqused_k,
+                            leftpad_k, block_table, alibi_slopes, max_seqlen_q,
+                            max_seqlen_k, dropout_p, softmax_scale,
+                            zero_tensors, causal, window_size_left,
+                            window_size_right, softcap, return_softmax,
+                            generator, flash_attn_mars_ext);
+  }
+};
+
+}  // namespace infini::ops::linked::torch::metax
+
+namespace infini::ops {
+
+void Operator<FlashAttnVarlenFunc, Device::Type::kMetax, 16>::operator()(
+    const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q,
+    const Tensor cu_seqlens_k, const std::optional<Tensor> alibi_slopes,
+    const std::optional<Tensor> block_table, const int64_t max_seqlen_q,
+    const int64_t max_seqlen_k, const double dropout_p,
+    const std::optional<double> softmax_scale, const bool causal,
+    const std::vector<int64_t> window_size, const double softcap,
+    const bool deterministic, const bool return_attn_probs, Tensor out,
+    std::optional<Tensor> softmax_lse, std::optional<Tensor> s_dmask) const {
+  using Delegate = linked::torch::TorchFlashAttnVarlenFunc<
+      linked::torch::metax::FlashAttnVarlen>;
+  if (!delegate_) {
+    delegate_ = std::make_unique<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);
+  }
+  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

src/linked/torch/metax/ops/flash_attn_varlen_func/flash_attn.h

@@ -0,0 +1,35 @@
+#ifndef INFINI_OPS_LINKED_TORCH_METAX_OPS_FLASH_ATTN_VARLEN_FUNC_FLASH_ATTN_H_
+#define INFINI_OPS_LINKED_TORCH_METAX_OPS_FLASH_ATTN_VARLEN_FUNC_FLASH_ATTN_H_
+
+#include <memory>
+
+#include "base/flash_attn_varlen_func.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<FlashAttnVarlenFunc, Device::Type::kMetax, 16>
+    : 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<Tensor> alibi_slopes,
+                  const std::optional<Tensor> block_table,
+                  const int64_t max_seqlen_q, const int64_t max_seqlen_k,
+                  const double dropout_p,
+                  const std::optional<double> softmax_scale, const bool causal,
+                  const std::vector<int64_t> window_size, const double softcap,
+                  const bool deterministic, const bool return_attn_probs,
+                  Tensor out, std::optional<Tensor> softmax_lse,
+                  std::optional<Tensor> s_dmask) const override;
+
+ private:
+  mutable std::unique_ptr<FlashAttnVarlenFunc> delegate_;
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_LINKED_TORCH_METAX_OPS_FLASH_ATTN_VARLEN_FUNC_FLASH_ATTN_H_

src/linked/torch/metax/ops/flash_attn_varlen_func/flash_attn.yaml

@@ -0,0 +1,4 @@
+library: flash_attn
+required_symbols:
+  - >-
+    mha_varlen_fwd(at::Tensor&, at::Tensor const&, at::Tensor const&, std::optional<at::Tensor>&, at::Tensor const&, at::Tensor const&, std::optional<at::Tensor>&, std::optional<at::Tensor const>&, std::optional<at::Tensor>&, std::optional<at::Tensor>&, int, int, float, float, bool, bool, int, int, float, bool, std::optional<at::Generator>, std::optional<at::Tensor>&)

src/linked/torch/metax/ops/flash_attn_with_kvcache/flash_attn.cc

@@ -0,0 +1,110 @@
+#include "linked/torch/metax/ops/flash_attn_with_kvcache/flash_attn.h"
+
+#include "linked/torch/ops/flash_attn_with_kvcache.h"
+#include "torch/metax/c10.h"
+
+std::vector<at::Tensor> mha_fwd_kvcache(
+    at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache,
+    std::optional<const at::Tensor>& k, std::optional<const at::Tensor>& v,
+    std::optional<const at::Tensor>& cache_seqlens,
+    std::optional<const at::Tensor>& rotary_cos,
+    std::optional<const at::Tensor>& rotary_sin,
+    std::optional<const at::Tensor>& cache_batch_idx,
+    std::optional<const at::Tensor>& cache_leftpad,
+    std::optional<at::Tensor>& block_table,
+    std::optional<at::Tensor>& alibi_slopes, std::optional<at::Tensor>& out,
+    float softmax_scale, bool causal, int window_size_left,
+    int window_size_right, float softcap, bool rotary_interleaved,
+    int num_splits, std::optional<at::Tensor>& flash_attn_mars_ext);
+
+namespace infini::ops::linked::torch::metax {
+
+struct FlashAttnKvcache : C10<Device::Type::kMetax> {
+  static std::vector<at::Tensor> Call(
+      at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache,
+      std::optional<const at::Tensor>& k, std::optional<const at::Tensor>& v,
+      std::optional<const at::Tensor>& cache_seqlens,
+      std::optional<const at::Tensor>& rotary_cos,
+      std::optional<const at::Tensor>& rotary_sin,
+      std::optional<const at::Tensor>& cache_batch_idx,
+      std::optional<const at::Tensor>& cache_leftpad,
+      std::optional<at::Tensor>& block_table,
+      std::optional<at::Tensor>& alibi_slopes, std::optional<at::Tensor>& out,
+      float softmax_scale, bool causal, int window_size_left,
+      int window_size_right, float softcap, bool rotary_interleaved,
+      int num_splits) {
+    std::optional<at::Tensor> flash_attn_mars_ext;
+    return ::mha_fwd_kvcache(
+        q, k_cache, v_cache, k, v, cache_seqlens, rotary_cos, rotary_sin,
+        cache_batch_idx, cache_leftpad, block_table, alibi_slopes, out,
+        softmax_scale, causal, window_size_left, window_size_right, softcap,
+        rotary_interleaved, num_splits, flash_attn_mars_ext);
+  }
+};
+
+}  // namespace infini::ops::linked::torch::metax
+
+namespace infini::ops {
+
+void Operator<FlashAttnWithKvcache, Device::Type::kMetax, 16>::operator()(
+    const Tensor q, Tensor k_cache, Tensor v_cache,
+    const std::optional<Tensor> k, const std::optional<Tensor> v,
+    const std::optional<Tensor> rotary_cos,
+    const std::optional<Tensor> rotary_sin, const int64_t cache_seqlens,
+    const std::optional<Tensor> cache_batch_idx,
+    const std::optional<Tensor> cache_leftpad,
+    const std::optional<Tensor> block_table,
+    const std::optional<Tensor> alibi_slopes,
+    const std::optional<double> softmax_scale, const bool causal,
+    const std::vector<int64_t> window_size, const double softcap,
+    const bool rotary_interleaved, const int64_t num_splits,
+    const bool return_softmax_lse, Tensor out,
+    std::optional<Tensor> softmax_lse) const {
+  using Delegate = linked::torch::TorchFlashAttnWithKvcache<
+      linked::torch::metax::FlashAttnKvcache>;
+  if (!delegate_) {
+    delegate_ = std::make_unique<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);
+  }
+  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<FlashAttnWithKvcache, Device::Type::kMetax, 16>::operator()(
+    const Tensor q, Tensor k_cache, Tensor v_cache,
+    const std::optional<Tensor> k, const std::optional<Tensor> v,
+    const std::optional<Tensor> rotary_cos,
+    const std::optional<Tensor> rotary_sin,
+    const std::optional<Tensor> cache_seqlens,
+    const std::optional<Tensor> cache_batch_idx,
+    const std::optional<Tensor> cache_leftpad,
+    const std::optional<Tensor> block_table,
+    const std::optional<Tensor> alibi_slopes,
+    const std::optional<double> softmax_scale, const bool causal,
+    const std::vector<int64_t> window_size, const double softcap,
+    const bool rotary_interleaved, const int64_t num_splits,
+    const bool return_softmax_lse, Tensor out,
+    std::optional<Tensor> softmax_lse) const {
+  using Delegate = linked::torch::TorchFlashAttnWithKvcache<
+      linked::torch::metax::FlashAttnKvcache>;
+  if (!delegate_) {
+    delegate_ = std::make_unique<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);
+  }
+  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

src/linked/torch/metax/ops/flash_attn_with_kvcache/flash_attn.h

@@ -0,0 +1,53 @@
+#ifndef INFINI_OPS_LINKED_TORCH_METAX_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_
+#define INFINI_OPS_LINKED_TORCH_METAX_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_
+
+#include <memory>
+
+#include "base/flash_attn_with_kvcache.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<FlashAttnWithKvcache, Device::Type::kMetax, 16>
+    : public FlashAttnWithKvcache {
+ public:
+  using FlashAttnWithKvcache::FlashAttnWithKvcache;
+  using FlashAttnWithKvcache::operator();
+
+  void operator()(const Tensor q, Tensor k_cache, Tensor v_cache,
+                  const std::optional<Tensor> k, const std::optional<Tensor> v,
+                  const std::optional<Tensor> rotary_cos,
+                  const std::optional<Tensor> rotary_sin,
+                  const int64_t cache_seqlens,
+                  const std::optional<Tensor> cache_batch_idx,
+                  const std::optional<Tensor> cache_leftpad,
+                  const std::optional<Tensor> block_table,
+                  const std::optional<Tensor> alibi_slopes,
+                  const std::optional<double> softmax_scale, const bool causal,
+                  const std::vector<int64_t> window_size, const double softcap,
+                  const bool rotary_interleaved, const int64_t num_splits,
+                  const bool return_softmax_lse, Tensor out,
+                  std::optional<Tensor> softmax_lse) const override;
+
+  void operator()(const Tensor q, Tensor k_cache, Tensor v_cache,
+                  const std::optional<Tensor> k, const std::optional<Tensor> v,
+                  const std::optional<Tensor> rotary_cos,
+                  const std::optional<Tensor> rotary_sin,
+                  const std::optional<Tensor> cache_seqlens,
+                  const std::optional<Tensor> cache_batch_idx,
+                  const std::optional<Tensor> cache_leftpad,
+                  const std::optional<Tensor> block_table,
+                  const std::optional<Tensor> alibi_slopes,
+                  const std::optional<double> softmax_scale, const bool causal,
+                  const std::vector<int64_t> window_size, const double softcap,
+                  const bool rotary_interleaved, const int64_t num_splits,
+                  const bool return_softmax_lse, Tensor out,
+                  std::optional<Tensor> softmax_lse) const override;
+
+ private:
+  mutable std::unique_ptr<FlashAttnWithKvcache> delegate_;
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_LINKED_TORCH_METAX_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_

src/linked/torch/metax/ops/flash_attn_with_kvcache/flash_attn.yaml

@@ -0,0 +1,4 @@
+library: flash_attn
+required_symbols:
+  - >-
+    mha_fwd_kvcache(at::Tensor&, at::Tensor const&, at::Tensor const&, std::optional<at::Tensor const>&, std::optional<at::Tensor const>&, std::optional<at::Tensor const>&, std::optional<at::Tensor const>&, std::optional<at::Tensor const>&, std::optional<at::Tensor const>&, std::optional<at::Tensor const>&, std::optional<at::Tensor>&, std::optional<at::Tensor>&, std::optional<at::Tensor>&, float, bool, int, int, float, bool, int, std::optional<at::Tensor>&)