From 9ea21d0dec64a5c1e4501a01f8adf2acb47831b5 Mon Sep 17 00:00:00 2001 From: Scott Wolchok Date: Fri, 10 May 2024 15:19:01 -0700 Subject: [PATCH] [executorch] Take Tensor for start_pos in llama custom SDPA Several different runtimes don't seem to like the custom SDPA taking a single scalar pulled from a Tensor as an argument: - dynamo doesn't support item() - CoreML doesn't seem to like this either So, just take a Tensor. Differential Revision: [D57233278](https://our.internmc.facebook.com/intern/diff/D57233278/) [ghstack-poisoned] --- examples/models/llama2/custom_ops/op_sdpa.cpp | 3 +- examples/models/llama2/custom_ops/op_sdpa.h | 2 +- .../models/llama2/custom_ops/op_sdpa_aot.cpp | 12 ++--- .../custom_ops/op_sdpa_with_kv_cache_test.cpp | 50 ++++++++++--------- .../llama2/source_transformation/sdpa.py | 2 +- 5 files changed, 36 insertions(+), 33 deletions(-) diff --git a/examples/models/llama2/custom_ops/op_sdpa.cpp b/examples/models/llama2/custom_ops/op_sdpa.cpp index 187ab08ffba..7e8d5a85fd0 100644 --- a/examples/models/llama2/custom_ops/op_sdpa.cpp +++ b/examples/models/llama2/custom_ops/op_sdpa.cpp @@ -732,7 +732,7 @@ Tensor& sdpa_with_kv_cache_out( const Tensor& v_projected, Tensor& key_cache, Tensor& value_cache, - const int64_t start_pos, + const Tensor& start_pos_tensor, const int64_t seq_len, const optional& attn_mask, const double dropout_p, @@ -741,6 +741,7 @@ Tensor& sdpa_with_kv_cache_out( const optional scale, Tensor& output) { (void)ctx; + const int64_t start_pos = *start_pos_tensor.const_data_ptr(); ET_KERNEL_CHECK( ctx, validate_cache_params(key_cache, value_cache, start_pos, seq_len), diff --git a/examples/models/llama2/custom_ops/op_sdpa.h b/examples/models/llama2/custom_ops/op_sdpa.h index fd130964ebb..ebcdac49501 100644 --- a/examples/models/llama2/custom_ops/op_sdpa.h +++ b/examples/models/llama2/custom_ops/op_sdpa.h @@ -22,7 +22,7 @@ Tensor& sdpa_with_kv_cache_out( const Tensor& v_projected, Tensor& key_cache, Tensor& value_cache, - const int64_t start_pos, + const Tensor& start_pos_tensor, const int64_t seq_len, const optional& attn_mask, const double dropout_p, diff --git a/examples/models/llama2/custom_ops/op_sdpa_aot.cpp b/examples/models/llama2/custom_ops/op_sdpa_aot.cpp index ed735406ad5..d084cafbe1e 100644 --- a/examples/models/llama2/custom_ops/op_sdpa_aot.cpp +++ b/examples/models/llama2/custom_ops/op_sdpa_aot.cpp @@ -23,7 +23,7 @@ Tensor& sdpa_with_kv_cache_out_no_context( const Tensor& v_projected, Tensor& key_cache, Tensor& value_cache, - const int64_t start_pos, + const Tensor& start_pos_tensor, const int64_t seq_len, // @lint-ignore CLANGTIDY facebook-hte-ConstantArgumentPassByValue // @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy @@ -41,7 +41,7 @@ Tensor& sdpa_with_kv_cache_out_no_context( v_projected, key_cache, value_cache, - start_pos, + start_pos_tensor, seq_len, attn_mask, dropout_p, @@ -56,7 +56,7 @@ at::Tensor sdpa_with_kv_cache_aten( const at::Tensor& v_projected, at::Tensor& key_cache, at::Tensor& value_cache, - const int64_t start_pos, + const at::Tensor& start_pos_tensor, const int64_t seq_len, // @lint-ignore CLANGTIDY facebook-hte-ConstantArgumentPassByValue // @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy @@ -72,7 +72,7 @@ at::Tensor sdpa_with_kv_cache_aten( v_projected, key_cache, value_cache, - start_pos, + start_pos_tensor, seq_len, attn_mask, dropout_p, @@ -89,11 +89,11 @@ at::Tensor sdpa_with_kv_cache_aten( TORCH_LIBRARY(llama, m) { m.def( "sdpa_with_kv_cache(Tensor query, Tensor key, Tensor value, Tensor(a!) key_cache, " - "Tensor(b!) value_cache, SymInt start_pos, SymInt seq_len, Tensor? attn_mask=None, " + "Tensor(b!) value_cache, Tensor start_pos, SymInt seq_len, Tensor? attn_mask=None, " "float drpout_p=0.0, bool is_causal=False, float? scale=None) -> Tensor"); m.def( "sdpa_with_kv_cache.out(Tensor query, Tensor key, Tensor value, Tensor(a!) key_cache, " - "Tensor(b!) value_cache, SymInt start_pos, SymInt seq_len, Tensor? attn_mask=None, " + "Tensor(b!) value_cache, Tensor start_pos, SymInt seq_len, Tensor? attn_mask=None, " "float drpout_p=0.0, bool is_causal=False, float? scale=None, *, Tensor(c!) out) -> Tensor(c!)"); } diff --git a/examples/models/llama2/custom_ops/op_sdpa_with_kv_cache_test.cpp b/examples/models/llama2/custom_ops/op_sdpa_with_kv_cache_test.cpp index fa2d164fe3d..f660b218d0d 100644 --- a/examples/models/llama2/custom_ops/op_sdpa_with_kv_cache_test.cpp +++ b/examples/models/llama2/custom_ops/op_sdpa_with_kv_cache_test.cpp @@ -18,13 +18,15 @@ using namespace ::testing; +using torch::executor::testing::TensorFactory; + exec_aten::Tensor op_sdpa_with_kv_cache( const exec_aten::Tensor& query, const exec_aten::Tensor& key, const exec_aten::Tensor& value, exec_aten::Tensor& key_cache, exec_aten::Tensor& value_cache, - const int64_t start_pos, + const exec_aten::Tensor& start_pos_tensor, const int64_t seq_len, const exec_aten::optional& attn_mask, double dropout_p, @@ -39,7 +41,7 @@ exec_aten::Tensor op_sdpa_with_kv_cache( value, key_cache, value_cache, - start_pos, + start_pos_tensor, seq_len, attn_mask, dropout_p, @@ -178,7 +180,7 @@ TEST(OpScaledDotProductAttentionTest, BasicTest) { value, key_cache_0, value_cache_0, - 0, + TensorFactory().make({1}, {0}), 1, attn_mask, dropout_p, @@ -213,7 +215,7 @@ TEST(OpScaledDotProductAttentionTest, BasicTest) { value, key_cache_2, value_cache_2, - 0, + TensorFactory().make({1}, {0}), 1, attn_mask, dropout_p, @@ -248,7 +250,7 @@ TEST(OpScaledDotProductAttentionTest, BasicTest) { value, key_cache_0, value_cache_0, - 1, + TensorFactory().make({1}, {1}), 1, attn_mask, dropout_p, @@ -283,7 +285,7 @@ TEST(OpScaledDotProductAttentionTest, BasicTest) { value, key_cache_1, value_cache_1, - 1, + TensorFactory().make({1}, {1}), 1, attn_mask, dropout_p, @@ -318,7 +320,7 @@ TEST(OpScaledDotProductAttentionTest, BasicTest) { value, key_cache_1, value_cache_1, - 2, + TensorFactory().make({1}, {2}), 1, attn_mask, dropout_p, @@ -353,7 +355,7 @@ TEST(OpScaledDotProductAttentionTest, BasicTest) { value, key_cache_2, value_cache_2, - 2, + TensorFactory().make({1}, {2}), 1, attn_mask, dropout_p, @@ -407,7 +409,7 @@ TEST(OpScaledDotProductAttentionTest, LargerTest) { value, key_cache_0, value_cache_0, - 0, + TensorFactory().make({1}, {0}), 1, attn_mask, dropout_p, @@ -429,7 +431,7 @@ TEST(OpScaledDotProductAttentionTest, LargerTest) { value, key_cache_2, value_cache_2, - 0, + TensorFactory().make({1}, {0}), 1, attn_mask, dropout_p, @@ -451,7 +453,7 @@ TEST(OpScaledDotProductAttentionTest, LargerTest) { value, key_cache_0, value_cache_0, - 1, + TensorFactory().make({1}, {1}), 1, attn_mask, dropout_p, @@ -473,7 +475,7 @@ TEST(OpScaledDotProductAttentionTest, LargerTest) { value, key_cache_1, value_cache_1, - 1, + TensorFactory().make({1}, {1}), 1, attn_mask, dropout_p, @@ -495,7 +497,7 @@ TEST(OpScaledDotProductAttentionTest, LargerTest) { value, key_cache_1, value_cache_1, - 2, + TensorFactory().make({1}, {2}), 1, attn_mask, dropout_p, @@ -517,7 +519,7 @@ TEST(OpScaledDotProductAttentionTest, LargerTest) { value, key_cache_2, value_cache_2, - 2, + TensorFactory().make({1}, {2}), 1, attn_mask, dropout_p, @@ -623,7 +625,7 @@ TEST(OpScaledDotProductAttentionTest, BasicTestWithAttnMask) { value, key_cache_0, value_cache_0, - 0, + TensorFactory().make({1}, {0}), 1, attn_mask, dropout_p, @@ -658,7 +660,7 @@ TEST(OpScaledDotProductAttentionTest, BasicTestWithAttnMask) { value, key_cache_2, value_cache_2, - 0, + TensorFactory().make({1}, {0}), 1, attn_mask, dropout_p, @@ -694,7 +696,7 @@ TEST(OpScaledDotProductAttentionTest, BasicTestWithAttnMask) { value, key_cache_0, value_cache_0, - 1, + TensorFactory().make({1}, {1}), 1, attn_mask, dropout_p, @@ -729,7 +731,7 @@ TEST(OpScaledDotProductAttentionTest, BasicTestWithAttnMask) { value, key_cache_1, value_cache_1, - 1, + TensorFactory().make({1}, {1}), 1, attn_mask, dropout_p, @@ -765,7 +767,7 @@ TEST(OpScaledDotProductAttentionTest, BasicTestWithAttnMask) { value, key_cache_1, value_cache_1, - 2, + TensorFactory().make({1}, {2}), 1, attn_mask, dropout_p, @@ -800,7 +802,7 @@ TEST(OpScaledDotProductAttentionTest, BasicTestWithAttnMask) { value, key_cache_2, value_cache_2, - 2, + TensorFactory().make({1}, {2}), 1, attn_mask, dropout_p, @@ -856,7 +858,7 @@ TEST(OpScaledDotProductAttentionTest, SequenceTest) { value, key_cache_0, value_cache_0, - 0, + TensorFactory().make({1}, {0}), 1, attn_mask, dropout_p, @@ -897,7 +899,7 @@ TEST(OpScaledDotProductAttentionTest, SequenceTest) { value, key_cache_0, value_cache_0, - 1, + TensorFactory().make({1}, {1}), 1, attn_mask, dropout_p, @@ -938,7 +940,7 @@ TEST(OpScaledDotProductAttentionTest, SequenceTest) { value, key_cache_0, value_cache_0, - 2, + TensorFactory().make({1}, {2}), 1, attn_mask, dropout_p, @@ -979,7 +981,7 @@ TEST(OpScaledDotProductAttentionTest, SequenceTest) { value, key_cache_0, value_cache_0, - 3, + TensorFactory().make({1}, {3}), 1, attn_mask, dropout_p, diff --git a/examples/models/llama2/source_transformation/sdpa.py b/examples/models/llama2/source_transformation/sdpa.py index 8a8a0cac7c2..897b95884f8 100644 --- a/examples/models/llama2/source_transformation/sdpa.py +++ b/examples/models/llama2/source_transformation/sdpa.py @@ -39,7 +39,7 @@ def forward( v, self.kv_cache.k_cache, self.kv_cache.v_cache, - input_pos[-1].item(), + input_pos[-1], seqlen, ) return output.view(bsz, seqlen, self.dim)