feat(thead): support FlashAttention on PPU

src/linked/torch/thead/flash_attn.yaml

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

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

@@ -0,0 +1,44 @@
+#include "linked/torch/thead/ops/flash_attn_varlen_func/flash_attn.h"
+
+namespace flash {
+
+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);
+
+}  // namespace flash
+
+namespace infini::ops::linked::torch::thead {
+
+std::vector<at::Tensor> FlashAttnVarlen::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) {
+  return flash::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);
+}
+
+}  // namespace infini::ops::linked::torch::thead
+
+namespace infini::ops::linked::torch {
+
+template class TorchFlashAttnVarlenFunc<
+    ::infini::ops::linked::torch::thead::FlashAttnVarlen>;
+
+}  // namespace infini::ops::linked::torch

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

@@ -0,0 +1,49 @@
+#ifndef INFINI_OPS_LINKED_TORCH_THEAD_OPS_FLASH_ATTN_VARLEN_FUNC_FLASH_ATTN_H_
+#define INFINI_OPS_LINKED_TORCH_THEAD_OPS_FLASH_ATTN_VARLEN_FUNC_FLASH_ATTN_H_
+
+#include <ATen/core/Generator.h>
+
+#include "linked/torch/ops/flash_attn_varlen_func.h"
+#include "torch/thead/c10.h"
+
+namespace infini::ops::linked::torch::thead {
+
+struct FlashAttnVarlen : C10<Device::Type::kThead> {
+  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);
+};
+
+}  // namespace infini::ops::linked::torch::thead
+
+namespace infini::ops::linked::torch {
+
+extern template class TorchFlashAttnVarlenFunc<
+    ::infini::ops::linked::torch::thead::FlashAttnVarlen>;
+
+}  // namespace infini::ops::linked::torch
+
+namespace infini::ops {
+
+template <>
+class Operator<FlashAttnVarlenFunc, Device::Type::kThead, 16>
+    : public linked::torch::TorchFlashAttnVarlenFunc<
+          linked::torch::thead::FlashAttnVarlen> {
+ public:
+  using linked::torch::TorchFlashAttnVarlenFunc<
+      linked::torch::thead::FlashAttnVarlen>::TorchFlashAttnVarlenFunc;
+
+  using linked::torch::TorchFlashAttnVarlenFunc<
+      linked::torch::thead::FlashAttnVarlen>::operator();
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_LINKED_TORCH_THEAD_OPS_FLASH_ATTN_VARLEN_FUNC_FLASH_ATTN_H_

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

@@ -0,0 +1,4 @@
+library: flash_attn
+required_symbols:
+  - >-
+    flash::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>)

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

@@ -0,0 +1,50 @@
+#include "linked/torch/thead/ops/flash_attn_with_kvcache/flash_attn.h"
+
+namespace flash {
+
+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);
+
+}  // namespace flash
+
+namespace infini::ops::linked::torch::thead {
+
+std::vector<at::Tensor> FlashAttnKvcache::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) {
+  return flash::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);
+}
+
+}  // namespace infini::ops::linked::torch::thead
+
+namespace infini::ops::linked::torch {
+
+template class TorchFlashAttnWithKvcache<
+    ::infini::ops::linked::torch::thead::FlashAttnKvcache>;
+
+}  // namespace infini::ops::linked::torch

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

@@ -0,0 +1,50 @@
+#ifndef INFINI_OPS_LINKED_TORCH_THEAD_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_
+#define INFINI_OPS_LINKED_TORCH_THEAD_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_
+
+#include "linked/torch/ops/flash_attn_with_kvcache.h"
+#include "torch/thead/c10.h"
+
+namespace infini::ops::linked::torch::thead {
+
+struct FlashAttnKvcache : C10<Device::Type::kThead> {
+  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);
+};
+
+}  // namespace infini::ops::linked::torch::thead
+
+namespace infini::ops::linked::torch {
+
+extern template class TorchFlashAttnWithKvcache<
+    ::infini::ops::linked::torch::thead::FlashAttnKvcache>;
+
+}  // namespace infini::ops::linked::torch
+
+namespace infini::ops {
+
+template <>
+class Operator<FlashAttnWithKvcache, Device::Type::kThead, 16>
+    : public linked::torch::TorchFlashAttnWithKvcache<
+          linked::torch::thead::FlashAttnKvcache> {
+ public:
+  using linked::torch::TorchFlashAttnWithKvcache<
+      linked::torch::thead::FlashAttnKvcache>::TorchFlashAttnWithKvcache;
+
+  using linked::torch::TorchFlashAttnWithKvcache<
+      linked::torch::thead::FlashAttnKvcache>::operator();
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_LINKED_TORCH_THEAD_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_

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

@@ -0,0 +1,4 @@
+library: flash_attn
+required_symbols:
+  - >-
+    flash::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)