Skip to content

Commit fa3f8ea

Browse files
voltjiavoltjia
committed
feat(linked): add MetaX FlashAttention providers
1 parent 8773433 commit fa3f8ea

7 files changed

Lines changed: 280 additions & 0 deletions

File tree

‎src/linked/torch/metax/flash_attn.yaml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
python_distribution_package: flash-attn
2+
library_glob: flash_attn_2_cuda*.so

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

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
#include "linked/torch/metax/ops/flash_attn_varlen_func/flash_attn.h"
2+
3+
#include <ATen/core/Generator.h>
4+
5+
#include "linked/torch/ops/flash_attn_varlen_func.h"
6+
#include "torch/metax/c10.h"
7+
8+
std::vector<at::Tensor> mha_varlen_fwd(
9+
at::Tensor& q, const at::Tensor& k, const at::Tensor& v,
10+
std::optional<at::Tensor>& out, const at::Tensor& cu_seqlens_q,
11+
const at::Tensor& cu_seqlens_k, std::optional<at::Tensor>& seqused_k,
12+
std::optional<const at::Tensor>& leftpad_k,
13+
std::optional<at::Tensor>& block_table,
14+
std::optional<at::Tensor>& alibi_slopes, int max_seqlen_q, int max_seqlen_k,
15+
float dropout_p, float softmax_scale, bool zero_tensors, bool causal,
16+
int window_size_left, int window_size_right, float softcap,
17+
bool return_softmax, std::optional<at::Generator> generator,
18+
std::optional<at::Tensor>& flash_attn_mars_ext);
19+
20+
namespace infini::ops::linked::torch::metax {
21+
22+
struct FlashAttnVarlen : C10<Device::Type::kMetax> {
23+
static std::vector<at::Tensor> Call(
24+
at::Tensor& q, const at::Tensor& k, const at::Tensor& v,
25+
std::optional<at::Tensor>& out, const at::Tensor& cu_seqlens_q,
26+
const at::Tensor& cu_seqlens_k, std::optional<at::Tensor>& seqused_k,
27+
std::optional<const at::Tensor>& leftpad_k,
28+
std::optional<at::Tensor>& block_table,
29+
std::optional<at::Tensor>& alibi_slopes, int max_seqlen_q,
30+
int max_seqlen_k, float dropout_p, float softmax_scale, bool zero_tensors,
31+
bool causal, int window_size_left, int window_size_right, float softcap,
32+
bool return_softmax, std::optional<at::Generator> generator) {
33+
std::optional<at::Tensor> flash_attn_mars_ext;
34+
return ::mha_varlen_fwd(q, k, v, out, cu_seqlens_q, cu_seqlens_k, seqused_k,
35+
leftpad_k, block_table, alibi_slopes, max_seqlen_q,
36+
max_seqlen_k, dropout_p, softmax_scale,
37+
zero_tensors, causal, window_size_left,
38+
window_size_right, softcap, return_softmax,
39+
generator, flash_attn_mars_ext);
40+
}
41+
};
42+
43+
} // namespace infini::ops::linked::torch::metax
44+
45+
namespace infini::ops {
46+
47+
void Operator<FlashAttnVarlenFunc, Device::Type::kMetax, 16>::operator()(
48+
const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q,
49+
const Tensor cu_seqlens_k, const std::optional<Tensor> alibi_slopes,
50+
const std::optional<Tensor> block_table, const int64_t max_seqlen_q,
51+
const int64_t max_seqlen_k, const double dropout_p,
52+
const std::optional<double> softmax_scale, const bool causal,
53+
const std::vector<int64_t> window_size, const double softcap,
54+
const bool deterministic, const bool return_attn_probs, Tensor out,
55+
std::optional<Tensor> softmax_lse, std::optional<Tensor> s_dmask) const {
56+
using Delegate = linked::torch::TorchFlashAttnVarlenFunc<
57+
linked::torch::metax::FlashAttnVarlen>;
58+
if (!delegate_) {
59+
delegate_ = std::make_unique<Delegate>(
60+
q, k, v, cu_seqlens_q, cu_seqlens_k, alibi_slopes, block_table,
61+
max_seqlen_q, max_seqlen_k, dropout_p, softmax_scale, causal,
62+
window_size, softcap, deterministic, return_attn_probs, out,
63+
softmax_lse, s_dmask);
64+
}
65+
delegate_->set_stream(stream_);
66+
(*delegate_)(q, k, v, cu_seqlens_q, cu_seqlens_k, alibi_slopes, block_table,
67+
max_seqlen_q, max_seqlen_k, dropout_p, softmax_scale, causal,
68+
window_size, softcap, deterministic, return_attn_probs, out,
69+
softmax_lse, s_dmask);
70+
}
71+
72+
} // namespace infini::ops

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
#ifndef INFINI_OPS_LINKED_TORCH_METAX_OPS_FLASH_ATTN_VARLEN_FUNC_FLASH_ATTN_H_
2+
#define INFINI_OPS_LINKED_TORCH_METAX_OPS_FLASH_ATTN_VARLEN_FUNC_FLASH_ATTN_H_
3+
4+
#include <memory>
5+
6+
#include "base/flash_attn_varlen_func.h"
7+
8+
namespace infini::ops {
9+
10+
template <>
11+
class Operator<FlashAttnVarlenFunc, Device::Type::kMetax, 16>
12+
: public FlashAttnVarlenFunc {
13+
public:
14+
using FlashAttnVarlenFunc::FlashAttnVarlenFunc;
15+
using FlashAttnVarlenFunc::operator();
16+
17+
void operator()(const Tensor q, const Tensor k, const Tensor v,
18+
const Tensor cu_seqlens_q, const Tensor cu_seqlens_k,
19+
const std::optional<Tensor> alibi_slopes,
20+
const std::optional<Tensor> block_table,
21+
const int64_t max_seqlen_q, const int64_t max_seqlen_k,
22+
const double dropout_p,
23+
const std::optional<double> softmax_scale, const bool causal,
24+
const std::vector<int64_t> window_size, const double softcap,
25+
const bool deterministic, const bool return_attn_probs,
26+
Tensor out, std::optional<Tensor> softmax_lse,
27+
std::optional<Tensor> s_dmask) const override;
28+
29+
private:
30+
mutable std::unique_ptr<FlashAttnVarlenFunc> delegate_;
31+
};
32+
33+
} // namespace infini::ops
34+
35+
#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‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
library: flash_attn
2+
required_symbols:
3+
- >-
4+
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‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
#include "linked/torch/metax/ops/flash_attn_with_kvcache/flash_attn.h"
2+
3+
#include "linked/torch/ops/flash_attn_with_kvcache.h"
4+
#include "torch/metax/c10.h"
5+
6+
std::vector<at::Tensor> mha_fwd_kvcache(
7+
at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache,
8+
std::optional<const at::Tensor>& k, std::optional<const at::Tensor>& v,
9+
std::optional<const at::Tensor>& cache_seqlens,
10+
std::optional<const at::Tensor>& rotary_cos,
11+
std::optional<const at::Tensor>& rotary_sin,
12+
std::optional<const at::Tensor>& cache_batch_idx,
13+
std::optional<const at::Tensor>& cache_leftpad,
14+
std::optional<at::Tensor>& block_table,
15+
std::optional<at::Tensor>& alibi_slopes, std::optional<at::Tensor>& out,
16+
float softmax_scale, bool causal, int window_size_left,
17+
int window_size_right, float softcap, bool rotary_interleaved,
18+
int num_splits, std::optional<at::Tensor>& flash_attn_mars_ext);
19+
20+
namespace infini::ops::linked::torch::metax {
21+
22+
struct FlashAttnKvcache : C10<Device::Type::kMetax> {
23+
static std::vector<at::Tensor> Call(
24+
at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache,
25+
std::optional<const at::Tensor>& k, std::optional<const at::Tensor>& v,
26+
std::optional<const at::Tensor>& cache_seqlens,
27+
std::optional<const at::Tensor>& rotary_cos,
28+
std::optional<const at::Tensor>& rotary_sin,
29+
std::optional<const at::Tensor>& cache_batch_idx,
30+
std::optional<const at::Tensor>& cache_leftpad,
31+
std::optional<at::Tensor>& block_table,
32+
std::optional<at::Tensor>& alibi_slopes, std::optional<at::Tensor>& out,
33+
float softmax_scale, bool causal, int window_size_left,
34+
int window_size_right, float softcap, bool rotary_interleaved,
35+
int num_splits) {
36+
std::optional<at::Tensor> flash_attn_mars_ext;
37+
return ::mha_fwd_kvcache(
38+
q, k_cache, v_cache, k, v, cache_seqlens, rotary_cos, rotary_sin,
39+
cache_batch_idx, cache_leftpad, block_table, alibi_slopes, out,
40+
softmax_scale, causal, window_size_left, window_size_right, softcap,
41+
rotary_interleaved, num_splits, flash_attn_mars_ext);
42+
}
43+
};
44+
45+
} // namespace infini::ops::linked::torch::metax
46+
47+
namespace infini::ops {
48+
49+
void Operator<FlashAttnWithKvcache, Device::Type::kMetax, 16>::operator()(
50+
const Tensor q, Tensor k_cache, Tensor v_cache,
51+
const std::optional<Tensor> k, const std::optional<Tensor> v,
52+
const std::optional<Tensor> rotary_cos,
53+
const std::optional<Tensor> rotary_sin, const int64_t cache_seqlens,
54+
const std::optional<Tensor> cache_batch_idx,
55+
const std::optional<Tensor> cache_leftpad,
56+
const std::optional<Tensor> block_table,
57+
const std::optional<Tensor> alibi_slopes,
58+
const std::optional<double> softmax_scale, const bool causal,
59+
const std::vector<int64_t> window_size, const double softcap,
60+
const bool rotary_interleaved, const int64_t num_splits,
61+
const bool return_softmax_lse, Tensor out,
62+
std::optional<Tensor> softmax_lse) const {
63+
using Delegate = linked::torch::TorchFlashAttnWithKvcache<
64+
linked::torch::metax::FlashAttnKvcache>;
65+
if (!delegate_) {
66+
delegate_ = std::make_unique<Delegate>(
67+
q, k_cache, v_cache, k, v, rotary_cos, rotary_sin, cache_seqlens,
68+
cache_batch_idx, cache_leftpad, block_table, alibi_slopes,
69+
softmax_scale, causal, window_size, softcap, rotary_interleaved,
70+
num_splits, return_softmax_lse, out, softmax_lse);
71+
}
72+
delegate_->set_stream(stream_);
73+
(*delegate_)(q, k_cache, v_cache, k, v, rotary_cos, rotary_sin, cache_seqlens,
74+
cache_batch_idx, cache_leftpad, block_table, alibi_slopes,
75+
softmax_scale, causal, window_size, softcap, rotary_interleaved,
76+
num_splits, return_softmax_lse, out, softmax_lse);
77+
}
78+
79+
void Operator<FlashAttnWithKvcache, Device::Type::kMetax, 16>::operator()(
80+
const Tensor q, Tensor k_cache, Tensor v_cache,
81+
const std::optional<Tensor> k, const std::optional<Tensor> v,
82+
const std::optional<Tensor> rotary_cos,
83+
const std::optional<Tensor> rotary_sin,
84+
const std::optional<Tensor> cache_seqlens,
85+
const std::optional<Tensor> cache_batch_idx,
86+
const std::optional<Tensor> cache_leftpad,
87+
const std::optional<Tensor> block_table,
88+
const std::optional<Tensor> alibi_slopes,
89+
const std::optional<double> softmax_scale, const bool causal,
90+
const std::vector<int64_t> window_size, const double softcap,
91+
const bool rotary_interleaved, const int64_t num_splits,
92+
const bool return_softmax_lse, Tensor out,
93+
std::optional<Tensor> softmax_lse) const {
94+
using Delegate = linked::torch::TorchFlashAttnWithKvcache<
95+
linked::torch::metax::FlashAttnKvcache>;
96+
if (!delegate_) {
97+
delegate_ = std::make_unique<Delegate>(
98+
q, k_cache, v_cache, k, v, rotary_cos, rotary_sin, cache_seqlens,
99+
cache_batch_idx, cache_leftpad, block_table, alibi_slopes,
100+
softmax_scale, causal, window_size, softcap, rotary_interleaved,
101+
num_splits, return_softmax_lse, out, softmax_lse);
102+
}
103+
delegate_->set_stream(stream_);
104+
(*delegate_)(q, k_cache, v_cache, k, v, rotary_cos, rotary_sin, cache_seqlens,
105+
cache_batch_idx, cache_leftpad, block_table, alibi_slopes,
106+
softmax_scale, causal, window_size, softcap, rotary_interleaved,
107+
num_splits, return_softmax_lse, out, softmax_lse);
108+
}
109+
110+
} // namespace infini::ops

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

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
#ifndef INFINI_OPS_LINKED_TORCH_METAX_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_
2+
#define INFINI_OPS_LINKED_TORCH_METAX_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_
3+
4+
#include <memory>
5+
6+
#include "base/flash_attn_with_kvcache.h"
7+
8+
namespace infini::ops {
9+
10+
template <>
11+
class Operator<FlashAttnWithKvcache, Device::Type::kMetax, 16>
12+
: public FlashAttnWithKvcache {
13+
public:
14+
using FlashAttnWithKvcache::FlashAttnWithKvcache;
15+
using FlashAttnWithKvcache::operator();
16+
17+
void operator()(const Tensor q, Tensor k_cache, Tensor v_cache,
18+
const std::optional<Tensor> k, const std::optional<Tensor> v,
19+
const std::optional<Tensor> rotary_cos,
20+
const std::optional<Tensor> rotary_sin,
21+
const int64_t cache_seqlens,
22+
const std::optional<Tensor> cache_batch_idx,
23+
const std::optional<Tensor> cache_leftpad,
24+
const std::optional<Tensor> block_table,
25+
const std::optional<Tensor> alibi_slopes,
26+
const std::optional<double> softmax_scale, const bool causal,
27+
const std::vector<int64_t> window_size, const double softcap,
28+
const bool rotary_interleaved, const int64_t num_splits,
29+
const bool return_softmax_lse, Tensor out,
30+
std::optional<Tensor> softmax_lse) const override;
31+
32+
void operator()(const Tensor q, Tensor k_cache, Tensor v_cache,
33+
const std::optional<Tensor> k, const std::optional<Tensor> v,
34+
const std::optional<Tensor> rotary_cos,
35+
const std::optional<Tensor> rotary_sin,
36+
const std::optional<Tensor> cache_seqlens,
37+
const std::optional<Tensor> cache_batch_idx,
38+
const std::optional<Tensor> cache_leftpad,
39+
const std::optional<Tensor> block_table,
40+
const std::optional<Tensor> alibi_slopes,
41+
const std::optional<double> softmax_scale, const bool causal,
42+
const std::vector<int64_t> window_size, const double softcap,
43+
const bool rotary_interleaved, const int64_t num_splits,
44+
const bool return_softmax_lse, Tensor out,
45+
std::optional<Tensor> softmax_lse) const override;
46+
47+
private:
48+
mutable std::unique_ptr<FlashAttnWithKvcache> delegate_;
49+
};
50+
51+
} // namespace infini::ops
52+
53+
#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‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
library: flash_attn
2+
required_symbols:
3+
- >-
4+
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>&)

0 commit comments

Comments
 (0)