Skip to content

Commit 804ebb5

Browse files
committed
Add contig requirement to aot binding for llm custom ops
Differential Revision: [D93870393](https://our.internmc.facebook.com/intern/diff/D93870393/) [ghstack-poisoned]
1 parent ffb29f7 commit 804ebb5

1 file changed

Lines changed: 21 additions & 10 deletions

File tree

extension/llm/custom_ops/op_sdpa_aot.cpp

Lines changed: 21 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -183,11 +183,11 @@ Tensor& sdpa_with_kv_cache_out_no_context(
183183
}
184184

185185
at::Tensor sdpa_with_kv_cache_aten(
186-
const at::Tensor& q_projected,
187-
const at::Tensor& k_projected,
188-
const at::Tensor& v_projected,
189-
at::Tensor& key_cache,
190-
at::Tensor& value_cache,
186+
const at::Tensor& q_proj,
187+
const at::Tensor& k_proj,
188+
const at::Tensor& v_proj,
189+
at::Tensor& k_cache,
190+
at::Tensor& v_cache,
191191
const int64_t start_pos,
192192
const int64_t seq_len,
193193
// @lint-ignore CLANGTIDY facebook-hte-ConstantArgumentPassByValue
@@ -197,6 +197,11 @@ at::Tensor sdpa_with_kv_cache_aten(
197197
const bool is_causal,
198198
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
199199
const std::optional<double> scale) {
200+
auto q_projected = q_proj.contiguous();
201+
auto k_projected = k_proj.contiguous();
202+
auto v_projected = v_proj.contiguous();
203+
auto key_cache = k_cache.contiguous();
204+
auto value_cache = v_cache.contiguous();
200205
auto output = at::empty_like(q_projected);
201206
WRAP_TO_ATEN(sdpa_with_kv_cache_out_no_context, 11)
202207
(q_projected,
@@ -256,11 +261,14 @@ at::Tensor custom_sdpa_aten(
256261
// @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
257262
const std::optional<double> scale,
258263
const bool is_seq_dim_2) {
259-
auto output = at::empty(q.sizes());
264+
auto q_projected = q.contiguous();
265+
auto k_projected = k.contiguous();
266+
auto v_projected = v.contiguous();
267+
auto output = at::empty_like(q_projected);
260268
WRAP_TO_ATEN(custom_sdpa_out_no_context, 9)
261-
(q,
262-
k,
263-
v,
269+
(q_projected,
270+
k_projected,
271+
v_projected,
264272
start_pos,
265273
attn_mask,
266274
dropout_p,
@@ -331,7 +339,10 @@ at::Tensor custom_quantized_sdpa_aten(
331339
const std::optional<at::Tensor>& v_zero_points,
332340
const std::optional<at::Tensor>& v_scales,
333341
const bool is_seq_at_dim_2) {
334-
auto output = at::empty(q.sizes());
342+
auto q_projected = q.contiguous();
343+
auto k_projected = k.contiguous();
344+
auto v_projected = v.contiguous();
345+
auto output = at::empty(q_projected.sizes());
335346
WRAP_TO_ATEN(custom_quantized_sdpa_out_no_context, 15)
336347
(q,
337348
k,

0 commit comments

Comments
 (0)