diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0402044da6..4a7c40d079 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2887,13 +2887,74 @@ public: for (auto & [buft, mbuf] : mbufs_new) { const auto & mbuf_cur = mbufs.at(buft); - if (!mbuf_cur.buf || mbuf_cur.n_tensors != mbuf.n_tensors || mbuf_cur.total_size != mbuf.total_size) { + if (!mbuf_cur.buf || mbuf_cur.total_size != mbuf.total_size) { GGML_ABORT("%s: memory buffer mismatch\n", __func__); } - for (size_t i = 0; i < mbuf_cur.org.size(); ++i) { - ggml_backend_tensor_copy(mbuf_cur.cpy[i], mbuf.org[i]); + // Copy the write-side data (mbuf_cur.cpy) into the read-side targets (mbuf.org) with a + // byte cursor. Write and read enumerate the same logical data in the same order but may + // chunk it differently, so copy across tensor boundaries rather than 1:1 by index. + const size_t total = mbuf_cur.total_size; + + ggml_init_params params_scratch = { + /*.mem_size =*/ 2*(mbuf_cur.cpy.size() + mbuf.org.size())*ggml_tensor_overhead(), + /*.mem_buffer =*/ NULL, + /*.no_alloc =*/ true, + }; + ggml_context * ctx_scratch = ggml_init(params_scratch); + + size_t src_pos = 0; + size_t dst_pos = 0; + size_t src_j = 0; + size_t dst_i = 0; + size_t src_base = 0; + size_t dst_base = 0; + + while (src_pos < total) { + const auto & src_t = mbuf_cur.cpy[src_j]; + const auto & dst_t = mbuf.org[dst_i]; + + const size_t src_size = ggml_nbytes(src_t); + const size_t dst_size = ggml_nbytes(dst_t); + + const size_t src_off = src_pos - src_base; + const size_t dst_off = dst_pos - dst_base; + + const size_t n_copy = std::min(src_size - src_off, dst_size - dst_off); + + const size_t el = ggml_element_size(src_t); + const int64_t n_el = (int64_t) (n_copy / el); + + auto * src_v = ggml_view_1d(ctx_scratch, src_t, n_el, src_off); + ggml_backend_view_init(src_v); + auto * dst_v = ggml_view_1d(ctx_scratch, dst_t, n_el, dst_off); + ggml_backend_view_init(dst_v); + + ggml_backend_tensor_copy(src_v, dst_v); + + src_pos += n_copy; + dst_pos += n_copy; + + if (src_pos - src_base == src_size) { + src_base = src_pos; + ++src_j; + } + if (dst_pos - dst_base == dst_size) { + dst_base = dst_pos; + ++dst_i; + } } + + GGML_ASSERT(src_pos == total && dst_pos == total); + // any tensors left unvisited hold no data + for (size_t i = src_j; i < mbuf_cur.cpy.size(); ++i) { + GGML_ASSERT(ggml_nbytes(mbuf_cur.cpy[i]) == 0); + } + for (size_t i = dst_i; i < mbuf.org.size(); ++i) { + GGML_ASSERT(ggml_nbytes(mbuf.org[i]) == 0); + } + + ggml_free(ctx_scratch); } GGML_ASSERT(buf_size == 0);