@@ -183,11 +183,11 @@ Tensor& sdpa_with_kv_cache_out_no_context(
183183}
184184
185185at::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