From 624fe512b484c65058ab2a913a466d5bdff7d7df Mon Sep 17 00:00:00 2001 From: zhushuang Date: Fri, 11 Sep 2026 11:08:08 +0800 Subject: [PATCH] feat(thead): support FlashAttention on PPU --- src/linked/torch/thead/flash_attn.yaml | 2 + .../ops/flash_attn_varlen_func/flash_attn.cc | 44 ++++++++++++++++ .../ops/flash_attn_varlen_func/flash_attn.h | 49 ++++++++++++++++++ .../flash_attn_varlen_func/flash_attn.yaml | 4 ++ .../ops/flash_attn_with_kvcache/flash_attn.cc | 50 +++++++++++++++++++ .../ops/flash_attn_with_kvcache/flash_attn.h | 50 +++++++++++++++++++ .../flash_attn_with_kvcache/flash_attn.yaml | 4 ++ 7 files changed, 203 insertions(+) create mode 100644 src/linked/torch/thead/flash_attn.yaml create mode 100644 src/linked/torch/thead/ops/flash_attn_varlen_func/flash_attn.cc create mode 100644 src/linked/torch/thead/ops/flash_attn_varlen_func/flash_attn.h create mode 100644 src/linked/torch/thead/ops/flash_attn_varlen_func/flash_attn.yaml create mode 100644 src/linked/torch/thead/ops/flash_attn_with_kvcache/flash_attn.cc create mode 100644 src/linked/torch/thead/ops/flash_attn_with_kvcache/flash_attn.h create mode 100644 src/linked/torch/thead/ops/flash_attn_with_kvcache/flash_attn.yaml diff --git a/src/linked/torch/thead/flash_attn.yaml b/src/linked/torch/thead/flash_attn.yaml new file mode 100644 index 000000000..c54d82dbd --- /dev/null +++ b/src/linked/torch/thead/flash_attn.yaml @@ -0,0 +1,2 @@ +python_distribution_package: flash-attn +library_glob: flash_attn_2_cuda*.so diff --git a/src/linked/torch/thead/ops/flash_attn_varlen_func/flash_attn.cc b/src/linked/torch/thead/ops/flash_attn_varlen_func/flash_attn.cc new file mode 100644 index 000000000..0be4d7491 --- /dev/null +++ b/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 mha_varlen_fwd( + at::Tensor& q, const at::Tensor& k, const at::Tensor& v, + std::optional& out, const at::Tensor& cu_seqlens_q, + const at::Tensor& cu_seqlens_k, std::optional& seqused_k, + std::optional& leftpad_k, + std::optional& block_table, + std::optional& 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 generator); + +} // namespace flash + +namespace infini::ops::linked::torch::thead { + +std::vector FlashAttnVarlen::Call( + at::Tensor& q, const at::Tensor& k, const at::Tensor& v, + std::optional& out, const at::Tensor& cu_seqlens_q, + const at::Tensor& cu_seqlens_k, std::optional& seqused_k, + std::optional& leftpad_k, + std::optional& block_table, + std::optional& 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 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 diff --git a/src/linked/torch/thead/ops/flash_attn_varlen_func/flash_attn.h b/src/linked/torch/thead/ops/flash_attn_varlen_func/flash_attn.h new file mode 100644 index 000000000..624de0548 --- /dev/null +++ b/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 + +#include "linked/torch/ops/flash_attn_varlen_func.h" +#include "torch/thead/c10.h" + +namespace infini::ops::linked::torch::thead { + +struct FlashAttnVarlen : C10 { + static std::vector Call( + at::Tensor& q, const at::Tensor& k, const at::Tensor& v, + std::optional& out, const at::Tensor& cu_seqlens_q, + const at::Tensor& cu_seqlens_k, std::optional& seqused_k, + std::optional& leftpad_k, + std::optional& block_table, + std::optional& 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 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 + : 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_ diff --git a/src/linked/torch/thead/ops/flash_attn_varlen_func/flash_attn.yaml b/src/linked/torch/thead/ops/flash_attn_varlen_func/flash_attn.yaml new file mode 100644 index 000000000..7f37cee62 --- /dev/null +++ b/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 const&, at::Tensor const&, std::optional&, std::optional&, std::optional&, std::optional&, int, int, float, float, bool, bool, int, int, float, bool, std::optional) diff --git a/src/linked/torch/thead/ops/flash_attn_with_kvcache/flash_attn.cc b/src/linked/torch/thead/ops/flash_attn_with_kvcache/flash_attn.cc new file mode 100644 index 000000000..d0e25c9a0 --- /dev/null +++ b/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 mha_fwd_kvcache( + at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& k, std::optional& v, + std::optional& cache_seqlens, + std::optional& rotary_cos, + std::optional& rotary_sin, + std::optional& cache_batch_idx, + std::optional& cache_leftpad, + std::optional& block_table, + std::optional& alibi_slopes, std::optional& 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 FlashAttnKvcache::Call( + at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& k, std::optional& v, + std::optional& cache_seqlens, + std::optional& rotary_cos, + std::optional& rotary_sin, + std::optional& cache_batch_idx, + std::optional& cache_leftpad, + std::optional& block_table, + std::optional& alibi_slopes, std::optional& 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 diff --git a/src/linked/torch/thead/ops/flash_attn_with_kvcache/flash_attn.h b/src/linked/torch/thead/ops/flash_attn_with_kvcache/flash_attn.h new file mode 100644 index 000000000..0c5b31cb8 --- /dev/null +++ b/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 { + static std::vector Call( + at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& k, std::optional& v, + std::optional& cache_seqlens, + std::optional& rotary_cos, + std::optional& rotary_sin, + std::optional& cache_batch_idx, + std::optional& cache_leftpad, + std::optional& block_table, + std::optional& alibi_slopes, std::optional& 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 + : 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_ diff --git a/src/linked/torch/thead/ops/flash_attn_with_kvcache/flash_attn.yaml b/src/linked/torch/thead/ops/flash_attn_with_kvcache/flash_attn.yaml new file mode 100644 index 000000000..f15198a47 --- /dev/null +++ b/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&, std::optional&, std::optional&, std::optional&, std::optional&, std::optional&, std::optional&, std::optional&, std::optional&, std::optional&, float, bool, int, int, float, bool, int)