Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions backends/cadence/fusion_g3/operators/op_dequantize.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@
#include <algorithm>
#include <cinttypes>
#include <cmath>
#include <optional>

#include <xa_nnlib_kernels_api.h>

Expand All @@ -23,6 +22,8 @@ using ::executorch::aten::ScalarType;
using ::executorch::aten::Tensor;
using ::executorch::runtime::KernelRuntimeContext;

template <typename T>
using optional = std::optional<T>;
/* ScalarType in Executorch do not have support for below data types.
* So, creating a placeholder for these data types. Once, ScalarTypes is
* updated to have support for below data types, these can be removed and
Expand Down Expand Up @@ -294,7 +295,7 @@ Tensor& dequantize_impl(
const auto* input_data_ptr = input.const_data_ptr<CTYPE_IN>(); \
ET_CHECK_MSG( \
*axis == 0, "Axis must be 0 for a single dimensional tensors"); \
const std::optional<int64_t> dim; \
const optional<int64_t> dim; \
torch::executor::apply_over_dim( \
[input_data_ptr, out_data_ptr, zero_point_data, scale_data]( \
size_t numel, size_t stride, size_t base_ix) { \
Expand Down Expand Up @@ -500,7 +501,7 @@ Tensor& dequantize_impl(
const auto* input_data_ptr = input.const_data_ptr<CTYPE_IN>(); \
ET_CHECK_MSG( \
*axis == 0, "Axis must be 0 for a single dimensional tensors"); \
const std::optional<int64_t> dim; \
const optional<int64_t> dim; \
torch::executor::apply_over_dim( \
[input_data_ptr, out_data_ptr, zero_point_data, scale_data]( \
size_t numel, size_t stride, size_t base_ix) { \
Expand Down
82 changes: 41 additions & 41 deletions extension/llm/custom_ops/op_sdpa.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ bool validate_flash_attention_args(
const Tensor& query,
const Tensor& key,
const Tensor& value,
const std::optional<Tensor>& attn_mask) {
const optional<Tensor>& attn_mask) {
ET_CHECK_OR_RETURN_FALSE(query.dim() == 4, "query must be a 4D tensor");
ET_CHECK_OR_RETURN_FALSE(key.dim() == 4, "key must be a 4D tensor");
ET_CHECK_OR_RETURN_FALSE(value.dim() == 4, "value must be a 4D tensor");
Expand Down Expand Up @@ -325,11 +325,11 @@ Tensor& flash_attention_kernel_out(
const Tensor& query,
const Tensor& key,
const Tensor& value,
const std::optional<Tensor>& attn_mask,
const optional<Tensor>& attn_mask,
const double dropout_p,
const bool is_causal,
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<double> scale,
const optional<double> scale,
Tensor& output) {
(void)ctx;
ET_KERNEL_CHECK(
Expand Down Expand Up @@ -362,12 +362,12 @@ Tensor& flash_attention_kernel_out(
is_causal,
attn_mask,
scale,
std::nullopt,
std::nullopt,
std::nullopt,
std::nullopt,
std::nullopt,
std::nullopt);
nullopt,
nullopt,
nullopt,
nullopt,
nullopt,
nullopt);
} else if (seq_len >= 192) {
sdpa::impl::cpu_flash_attention<CTYPE, 64, 512>(
ctx,
Expand All @@ -379,12 +379,12 @@ Tensor& flash_attention_kernel_out(
is_causal,
attn_mask,
scale,
std::nullopt,
std::nullopt,
std::nullopt,
std::nullopt,
std::nullopt,
std::nullopt);
nullopt,
nullopt,
nullopt,
nullopt,
nullopt,
nullopt);
} else {
sdpa::impl::cpu_flash_attention<CTYPE, 32, 512>(
ctx,
Expand All @@ -396,12 +396,12 @@ Tensor& flash_attention_kernel_out(
is_causal,
attn_mask,
scale,
std::nullopt,
std::nullopt,
std::nullopt,
std::nullopt,
std::nullopt,
std::nullopt);
nullopt,
nullopt,
nullopt,
nullopt,
nullopt,
nullopt);
}
});
return output;
Expand All @@ -413,18 +413,18 @@ Tensor& custom_sdpa_out_impl(
const Tensor& k,
const Tensor& v,
const int64_t start_pos,
const std::optional<Tensor>& attn_mask,
const optional<Tensor>& attn_mask,
const double dropout_p,
const bool is_causal,
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<double> scale,
const optional<double> scale,
Tensor& output,
const std::optional<Tensor>& q_zero_points = std::nullopt,
const std::optional<Tensor>& q_scales = std::nullopt,
const std::optional<Tensor>& k_zero_points = std::nullopt,
const std::optional<Tensor>& k_scales = std::nullopt,
const std::optional<Tensor>& v_zero_points = std::nullopt,
const std::optional<Tensor>& v_scales = std::nullopt,
const optional<Tensor>& q_zero_points = nullopt,
const optional<Tensor>& q_scales = nullopt,
const optional<Tensor>& k_zero_points = nullopt,
const optional<Tensor>& k_scales = nullopt,
const optional<Tensor>& v_zero_points = nullopt,
const optional<Tensor>& v_scales = nullopt,
bool is_seq_at_dim_2 = false) {
ET_KERNEL_CHECK_MSG(
ctx,
Expand Down Expand Up @@ -570,17 +570,17 @@ Tensor& custom_quantized_sdpa_out(
const Tensor& k,
const Tensor& v,
const int64_t start_pos,
const std::optional<Tensor>& attn_mask,
const optional<Tensor>& attn_mask,
const double dropout_p,
const bool is_causal,
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<double> scale,
const std::optional<Tensor>& q_zero_points,
const std::optional<Tensor>& q_scales,
const std::optional<Tensor>& k_zero_points,
const std::optional<Tensor>& k_scales,
const std::optional<Tensor>& v_zero_points,
const std::optional<Tensor>& v_scales,
const optional<double> scale,
const optional<Tensor>& q_zero_points,
const optional<Tensor>& q_scales,
const optional<Tensor>& k_zero_points,
const optional<Tensor>& k_scales,
const optional<Tensor>& v_zero_points,
const optional<Tensor>& v_scales,
const bool is_seq_at_dim_2,
Tensor& output) {
return custom_sdpa_out_impl(
Expand Down Expand Up @@ -624,11 +624,11 @@ Tensor& custom_sdpa_out(
const Tensor& k,
const Tensor& v,
const int64_t start_pos,
const std::optional<Tensor>& attn_mask,
const optional<Tensor>& attn_mask,
const double dropout_p,
const bool is_causal,
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<double> scale,
const optional<double> scale,
Tensor& output) {
return custom_sdpa_out_impl(
ctx, q, k, v, start_pos, attn_mask, dropout_p, is_causal, scale, output);
Expand Down Expand Up @@ -658,11 +658,11 @@ Tensor& sdpa_with_kv_cache_out(
Tensor& value_cache,
const int64_t start_pos,
const int64_t seq_len,
const std::optional<Tensor>& attn_mask,
const optional<Tensor>& attn_mask,
const double dropout_p,
const bool is_causal,
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<double> scale,
const optional<double> scale,
Tensor& output) {
(void)ctx;
ET_KERNEL_CHECK(
Expand Down
28 changes: 14 additions & 14 deletions extension/llm/custom_ops/op_sdpa.h
Original file line number Diff line number Diff line change
Expand Up @@ -26,11 +26,11 @@ Tensor& sdpa_with_kv_cache_out(
Tensor& value_cache,
const int64_t start_pos,
const int64_t seq_len,
const std::optional<Tensor>& attn_mask,
const optional<Tensor>& attn_mask,
const double dropout_p,
const bool is_causal,
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<double> scale,
const optional<double> scale,
Tensor& output);

Tensor& custom_sdpa_out(
Expand All @@ -39,23 +39,23 @@ Tensor& custom_sdpa_out(
const Tensor& k,
const Tensor& v,
const int64_t start_pos,
const std::optional<Tensor>& attn_mask,
const optional<Tensor>& attn_mask,
const double dropout_p,
const bool is_causal,
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<double> scale,
const optional<double> scale,
Tensor& output);

Tensor& flash_attention_kernel_out(
KernelRuntimeContext& ctx,
const Tensor& query,
const Tensor& key,
const Tensor& value,
const std::optional<Tensor>& attn_mask,
const optional<Tensor>& attn_mask,
const double dropout_p,
const bool is_causal,
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<double> scale,
const optional<double> scale,
Tensor& output);

Tensor& custom_quantized_sdpa_out(
Expand All @@ -64,17 +64,17 @@ Tensor& custom_quantized_sdpa_out(
const Tensor& k,
const Tensor& v,
const int64_t start_pos,
const std::optional<Tensor>& attn_mask,
const optional<Tensor>& attn_mask,
const double dropout_p,
const bool is_causal,
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<double> scale,
const std::optional<Tensor>& q_zero_points,
const std::optional<Tensor>& q_scales,
const std::optional<Tensor>& k_zero_points,
const std::optional<Tensor>& k_scales,
const std::optional<Tensor>& v_zero_points,
const std::optional<Tensor>& v_scales,
const optional<double> scale,
const optional<Tensor>& q_zero_points,
const optional<Tensor>& q_scales,
const optional<Tensor>& k_zero_points,
const optional<Tensor>& k_scales,
const optional<Tensor>& v_zero_points,
const optional<Tensor>& v_scales,
const bool is_seq_at_dim_1,
Tensor& output);

Expand Down
48 changes: 24 additions & 24 deletions extension/llm/custom_ops/op_sdpa_aot.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -27,11 +27,11 @@ Tensor& sdpa_with_kv_cache_out_no_context(
const int64_t seq_len,
// @lint-ignore CLANGTIDY facebook-hte-ConstantArgumentPassByValue
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<Tensor> attn_mask,
const optional<Tensor> attn_mask,
const double dropout_p,
const bool is_causal,
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<double> scale,
const optional<double> scale,
Tensor& output);

at::Tensor sdpa_with_kv_cache_aten(
Expand All @@ -57,11 +57,11 @@ Tensor& custom_sdpa_out_no_context(
const int64_t start_pos,
// @lint-ignore CLANGTIDY facebook-hte-ConstantArgumentPassByValue
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<Tensor> attn_mask,
const optional<Tensor> attn_mask,
const double dropout_p,
const bool is_causal,
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<double> scale,
const optional<double> scale,
Tensor& output);

at::Tensor custom_sdpa_aten(
Expand All @@ -84,17 +84,17 @@ Tensor& custom_quantized_sdpa_out_no_context(
const int64_t start_pos,
// @lint-ignore CLANGTIDY facebook-hte-ConstantArgumentPassByValue
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<Tensor> attn_mask,
const optional<Tensor> attn_mask,
const double dropout_p,
const bool is_causal,
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<double> scale,
const std::optional<Tensor> q_zero_points,
const std::optional<Tensor> q_scales,
const std::optional<Tensor> k_zero_points,
const std::optional<Tensor> k_scales,
const std::optional<Tensor> v_zero_points,
const std::optional<Tensor> v_scales,
const optional<double> scale,
const optional<Tensor> q_zero_points,
const optional<Tensor> q_scales,
const optional<Tensor> k_zero_points,
const optional<Tensor> k_scales,
const optional<Tensor> v_zero_points,
const optional<Tensor> v_scales,
const bool is_seq_at_dim_2,
Tensor& output);

Expand Down Expand Up @@ -153,11 +153,11 @@ Tensor& sdpa_with_kv_cache_out_no_context(
const int64_t seq_len,
// @lint-ignore CLANGTIDY facebook-hte-ConstantArgumentPassByValue
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<Tensor> attn_mask,
const optional<Tensor> attn_mask,
const double dropout_p,
const bool is_causal,
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<double> scale,
const optional<double> scale,
Tensor& output) {
executorch::runtime::KernelRuntimeContext context{};
return torch::executor::native::sdpa_with_kv_cache_out(
Expand Down Expand Up @@ -215,11 +215,11 @@ Tensor& custom_sdpa_out_no_context(
const int64_t start_pos,
// @lint-ignore CLANGTIDY facebook-hte-ConstantArgumentPassByValue
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<Tensor> attn_mask,
const optional<Tensor> attn_mask,
const double dropout_p,
const bool is_causal,
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<double> scale,
const optional<double> scale,
Tensor& output) {
executorch::aten::RuntimeContext context{};
return torch::executor::native::custom_sdpa_out(
Expand Down Expand Up @@ -260,17 +260,17 @@ Tensor& custom_quantized_sdpa_out_no_context(
const int64_t start_pos,
// @lint-ignore CLANGTIDY facebook-hte-ConstantArgumentPassByValue
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<Tensor> attn_mask,
const optional<Tensor> attn_mask,
const double dropout_p,
const bool is_causal,
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
const std::optional<double> scale,
const std::optional<Tensor> q_zero_points,
const std::optional<Tensor> q_scales,
const std::optional<Tensor> k_zero_points,
const std::optional<Tensor> k_scales,
const std::optional<Tensor> v_zero_points,
const std::optional<Tensor> v_scales,
const optional<double> scale,
const optional<Tensor> q_zero_points,
const optional<Tensor> q_scales,
const optional<Tensor> k_zero_points,
const optional<Tensor> k_scales,
const optional<Tensor> v_zero_points,
const optional<Tensor> v_scales,
const bool is_seq_at_dim_2,
Tensor& output) {
executorch::aten::RuntimeContext context{};
Expand Down
16 changes: 8 additions & 8 deletions extension/llm/custom_ops/op_sdpa_impl.h
Original file line number Diff line number Diff line change
Expand Up @@ -613,14 +613,14 @@ void cpu_flash_attention(
const Tensor& value,
double dropout_p,
bool is_causal,
const std::optional<Tensor>& attn_mask,
const std::optional<double>& scale,
const std::optional<Tensor>& q_zero_points,
const std::optional<Tensor>& q_scales,
const std::optional<Tensor>& k_zero_points,
const std::optional<Tensor>& k_scales,
const std::optional<Tensor>& v_zero_points,
const std::optional<Tensor>& v_scales,
const optional<Tensor>& attn_mask,
const optional<double>& scale,
const optional<Tensor>& q_zero_points,
const optional<Tensor>& q_scales,
const optional<Tensor>& k_zero_points,
const optional<Tensor>& k_scales,
const optional<Tensor>& v_zero_points,
const optional<Tensor>& v_scales,
const SeqDim seq_dim = SeqDim::TWO,
const int64_t start_pos = 0,
const int64_t num_keys_for_causal_attention = -1) {
Expand Down
Loading
Loading