Qualcomm AI Engine Direct - Fixed the wrong size for memmove in kv_manager (pytorch#14685)

shewu-quic · web-flow · commit f662cf5b29ab · 2025-09-30T09:53:39.000-07:00
As title. cc: @haowhsu-quic , @winskuo-quic
diff --git a/examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp b/examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp
@@ -365,7 +365,8 @@ void KVManager<T>::rearrange_key(KVCache<T>& k_cache, int32_t ar_len_dst) {
     }
     // copy from first dimension
     for (int i = 0; i < metadata_.head_dim; i++) {
-      std::memmove(k_cache_in_write_ptr, k_cache_in_read_ptr, dst_cache_num);
+      std::memmove(
+          k_cache_in_write_ptr, k_cache_in_read_ptr, dst_cache_num * sizeof(T));
       k_cache_in_read_ptr += src_cache_num;
       k_cache_in_write_ptr += dst_cache_num;
     }
@@ -378,7 +379,8 @@ void KVManager<T>::rearrange_key(KVCache<T>& k_cache, int32_t ar_len_dst) {
     }
     // copy from last dimension
     for (int i = 0; i < metadata_.head_dim; i++) {
-      std::memmove(k_cache_in_write_ptr, k_cache_in_read_ptr, src_cache_num);
+      std::memmove(
+          k_cache_in_write_ptr, k_cache_in_read_ptr, src_cache_num * sizeof(T));
       k_cache_in_read_ptr -= src_cache_num;
       k_cache_in_write_ptr -= dst_cache_num;
     }

Original file line number	Diff line number	Diff line change
`@@ -365,7 +365,8 @@ void KVManager<T>::rearrange_key(KVCache<T>& k_cache, int32_t ar_len_dst) {`
`365`	`365`	`}`
`366`	`366`	`// copy from first dimension`
`367`	`367`	`for (int i = 0; i < metadata_.head_dim; i++) {`
`368`		`- std::memmove(k_cache_in_write_ptr, k_cache_in_read_ptr, dst_cache_num);`
	`368`	`+ std::memmove(`
	`369`	`+ k_cache_in_write_ptr, k_cache_in_read_ptr, dst_cache_num * sizeof(T));`
`369`	`370`	`k_cache_in_read_ptr += src_cache_num;`
`370`	`371`	`k_cache_in_write_ptr += dst_cache_num;`
`371`	`372`	`}`
`@@ -378,7 +379,8 @@ void KVManager<T>::rearrange_key(KVCache<T>& k_cache, int32_t ar_len_dst) {`
`378`	`379`	`}`
`379`	`380`	`// copy from last dimension`
`380`	`381`	`for (int i = 0; i < metadata_.head_dim; i++) {`
`381`		`- std::memmove(k_cache_in_write_ptr, k_cache_in_read_ptr, src_cache_num);`
	`382`	`+ std::memmove(`
	`383`	`+ k_cache_in_write_ptr, k_cache_in_read_ptr, src_cache_num * sizeof(T));`
`382`	`384`	`k_cache_in_read_ptr -= src_cache_num;`
`383`	`385`	`k_cache_in_write_ptr -= dst_cache_num;`
`384`	`386`	`}`