feat(mars): support flash-attn in mars

CMakeLists.txt

@@ -240,7 +240,11 @@ if(AUTO_DETECT_BACKENDS)
 endif()
 
 if(WITH_TORCH OR WITH_LINKED)
-    find_package(Python COMPONENTS Interpreter REQUIRED)
+    if(WITH_MARS AND WITH_TORCH AND WITH_LINKED)
+        find_package(Python COMPONENTS Interpreter Development.Embed REQUIRED)
+    else()
+        find_package(Python COMPONENTS Interpreter REQUIRED)
+    endif()
 
     # Prefer the interpreter that the auto-detect block already
     # confirmed has `torch` (this is the system Python on hosts that
@@ -288,7 +292,7 @@ if(WITH_TORCH OR WITH_LINKED)
     find_library(C10_LIB c10 HINTS ${_torch_lib_dirs} REQUIRED)
     set(TORCH_LIBRARIES ${TORCH_LIB} ${TORCH_CPU_LIB} ${C10_LIB})
 
-    if(WITH_NVIDIA OR WITH_ILUVATAR OR WITH_THEAD)
+    if(WITH_NVIDIA OR WITH_ILUVATAR OR WITH_THEAD OR WITH_MARS)
         find_library(TORCH_CUDA_LIB torch_cuda HINTS ${_torch_lib_dirs} REQUIRED)
         find_library(C10_CUDA_LIB c10_cuda HINTS ${_torch_lib_dirs} REQUIRED)
         # `torch_cuda` registers CUDA kernels through static initializers, so
@@ -304,6 +308,12 @@ if(WITH_TORCH OR WITH_LINKED)
             -Wl,--no-as-needed ${TORCH_HIP_LIB} ${C10_HIP_LIB} -Wl,--as-needed)
     endif()
 
+    if(WITH_MARS AND WITH_TORCH AND WITH_LINKED)
+        # The Mars FlashAttention wheel loads libtorch_python without declaring
+        # its libpython dependency, so keep the embedding library in our chain.
+        list(APPEND TORCH_LIBRARIES Python::Python)
+    endif()
+
     if(WITH_METAX)
         find_library(C10_CUDA_LIB c10_cuda HINTS ${_torch_lib_dirs} REQUIRED)
         find_library(MACA_TORCH_RUNTIME_LIB runtime_cu

src/linked/torch/mars/flash_attn.yaml

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

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

@@ -0,0 +1,72 @@
+#include "linked/torch/mars/ops/flash_attn_varlen_func/flash_attn.h"
+
+#include <ATen/core/Generator.h>
+
+#include "linked/torch/ops/flash_attn_varlen_func.h"
+#include "torch/mars/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, bool return_max_logit);
+
+namespace infini::ops::linked::torch::mars {
+
+struct FlashAttnVarlen : C10<Device::Type::kMars> {
+  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, false);
+  }
+};
+
+}  // namespace infini::ops::linked::torch::mars
+
+namespace infini::ops {
+
+void Operator<FlashAttnVarlenFunc, Device::Type::kMars, 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::mars::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/mars/ops/flash_attn_varlen_func/flash_attn.h

@@ -0,0 +1,35 @@
+#ifndef INFINI_OPS_LINKED_TORCH_MARS_OPS_FLASH_ATTN_VARLEN_FUNC_FLASH_ATTN_H_
+#define INFINI_OPS_LINKED_TORCH_MARS_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::kMars, 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_MARS_OPS_FLASH_ATTN_VARLEN_FUNC_FLASH_ATTN_H_

src/linked/torch/mars/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>&, bool)

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

@@ -0,0 +1,110 @@
+#include "linked/torch/mars/ops/flash_attn_with_kvcache/flash_attn.h"
+
+#include "linked/torch/ops/flash_attn_with_kvcache.h"
+#include "torch/mars/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::mars {
+
+struct FlashAttnKvcache : C10<Device::Type::kMars> {
+  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::mars
+
+namespace infini::ops {
+
+void Operator<FlashAttnWithKvcache, Device::Type::kMars, 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::mars::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::kMars, 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::mars::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/mars/ops/flash_attn_with_kvcache/flash_attn.h

@@ -0,0 +1,53 @@
+#ifndef INFINI_OPS_LINKED_TORCH_MARS_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_
+#define INFINI_OPS_LINKED_TORCH_MARS_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::kMars, 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_MARS_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_

src/linked/torch/mars/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>&)

tests/conftest.py

@@ -25,7 +25,7 @@ def pytest_addoption(parser):
         "--devices",
         nargs="+",
         default=None,
-        help="Device(s) to test on (e.g., `--devices ascend cpu`). Accepts platform names (`nvidia`, `metax`, `iluvatar`, `hygon`, `thead`, `moore`, `cambricon`, `ascend`) or PyTorch device types (`cuda`, `mlu`, `musa`, `npu`). Defaults to all available devices.",
+        help="Device(s) to test on (e.g., `--devices ascend cpu`). Accepts platform names (`nvidia`, `mars`, `metax`, `iluvatar`, `hygon`, `thead`, `moore`, `cambricon`, `ascend`) or PyTorch device types (`cuda`, `mlu`, `musa`, `npu`). Defaults to all available devices.",
     )
     parser.addoption(
         "--report",
@@ -191,6 +191,7 @@ def _set_random_seed(seed):
 
 _PLATFORM_TO_TORCH_DEVICE = {
     "nvidia": "cuda",
+    "mars": "cuda",
     "metax": "cuda",
     "iluvatar": "cuda",
     "hygon": "cuda",