kompute : make partial tensor copies faster by syncing less data (#15)

Signed-off-by: Jared Van Bortel <[email protected]>

Author: cebtenzzre

ggml-kompute.cpp

@@ -62,6 +62,8 @@
 
 typedef ggml_fp16_t half;
 
+static const std::shared_ptr<kp::Tensor> nullTensor = nullptr;
+
 static std::string ggml_kompute_format_name(int device) {
     return "Kompute" + std::to_string(device);
 }
@@ -585,31 +587,47 @@ ggml_vk_memory * ggml_vk_find_tensor(const struct ggml_tensor * t, uint64_t & of
 }
 
 static
-const std::shared_ptr<kp::Tensor> ggml_vk_get_tensor(const struct ggml_tensor * t, uint32_t * alignedOffset = nullptr) {
-    uint64_t originalOffset = 0;
-    auto * res = ggml_vk_find_tensor(t, originalOffset);
+const std::shared_ptr<kp::Tensor> ggml_vk_get_tensor_aligned(const struct ggml_tensor * t, uint32_t * aligned_offset) {
+    uint64_t original_offset = 0;
+    auto * res = ggml_vk_find_tensor(t, original_offset);
     if (!res) {
-        static std::shared_ptr<kp::Tensor> nullTensor = nullptr;
         return nullTensor;
     }
 
     // Create a tensor whose memory will be composed of our buffers at the correct offset
-    const size_t nelements = ggml_nelements(t);
     size_t nbytes = ggml_nbytes(t);
+    size_t vulkan_offset = ggml_vk_aligned_offset(t->buffer, original_offset);
+    *aligned_offset = original_offset - vulkan_offset;
+    nbytes += *aligned_offset;
+
+    return komputeManager()->tensor(
+        t->data,
+        ggml_nelements(t), nbytes,
+        kp::Tensor::TensorDataTypes::eFloat,
+        res->primaryMemory, res->primaryBuffer,
+        res->stagingMemory, res->stagingBuffer,
+        vulkan_offset);
+}
 
-    size_t vulkanOffset = ggml_vk_aligned_offset(t->buffer, originalOffset);
-    if (alignedOffset) {
-        *alignedOffset = originalOffset - vulkanOffset;
-        nbytes += *alignedOffset;
+static
+const std::shared_ptr<kp::Tensor> ggml_vk_get_tensor_slice(const struct ggml_tensor * t, size_t offset, size_t nbytes) {
+    uint64_t tensor_offset = 0;
+    auto * res = ggml_vk_find_tensor(t, tensor_offset);
+    if (!res) {
+        return nullTensor;
     }
 
+    size_t elsz = ggml_element_size(t);
+    GGML_ASSERT(nbytes % elsz == 0);
+
+    // Create a tensor whose memory will be composed of our buffers at the correct offset
     return komputeManager()->tensor(
-        t->data,
-        nelements,
-        nbytes, kp::Tensor::TensorDataTypes::eFloat,
+        reinterpret_cast<char *>(t->data) + offset,
+        nbytes / elsz, nbytes,
+        kp::Tensor::TensorDataTypes::eFloat,
         res->primaryMemory, res->primaryBuffer,
         res->stagingMemory, res->stagingBuffer,
-        vulkanOffset);
+        tensor_offset + offset);
 }
 
 static std::vector<uint32_t> getSpirvShader(const unsigned char* rawData, size_t size) {
@@ -1551,13 +1569,12 @@ static void ggml_vk_graph_compute(struct ggml_kompute_context * ctx, struct ggml
             const enum ggml_type src1t = src1 ? src1->type : GGML_TYPE_COUNT;
             const enum ggml_type dstt = dst ? dst->type : GGML_TYPE_COUNT;
 
-            const static std::shared_ptr<kp::Tensor> nullTensor = nullptr;
             uint32_t off_src0 = 0;
             uint32_t off_src1 = 0;
             uint32_t off_dst  = 0;
-            const std::shared_ptr<kp::Tensor>& id_src0 = src0 ? ggml_vk_get_tensor(src0, &off_src0) : nullTensor;
-            const std::shared_ptr<kp::Tensor>& id_src1 = src1 ? ggml_vk_get_tensor(src1, &off_src1) : nullTensor;
-            const std::shared_ptr<kp::Tensor>& id_dst  = dst  ? ggml_vk_get_tensor(dst,  &off_dst)  : nullTensor;
+            const std::shared_ptr<kp::Tensor>& id_src0 = src0 ? ggml_vk_get_tensor_aligned(src0, &off_src0) : nullTensor;
+            const std::shared_ptr<kp::Tensor>& id_src1 = src1 ? ggml_vk_get_tensor_aligned(src1, &off_src1) : nullTensor;
+            const std::shared_ptr<kp::Tensor>& id_dst  = dst  ? ggml_vk_get_tensor_aligned(dst,  &off_dst)  : nullTensor;
 
             switch (dst->op) {
                 case GGML_OP_ADD:
@@ -1876,7 +1893,7 @@ static void * ggml_backend_kompute_buffer_get_base(ggml_backend_buffer_t buffer)
 static void ggml_backend_kompute_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
     GGML_UNUSED(buffer);
 
-    const auto res = ggml_vk_get_tensor(tensor);
+    const auto res = ggml_vk_get_tensor_slice(tensor, offset, size);
     GGML_ASSERT(res);
 
     memcpy((char *)tensor->data + offset, data, size);
@@ -1887,7 +1904,7 @@ static void ggml_backend_kompute_buffer_set_tensor(ggml_backend_buffer_t buffer,
 static void ggml_backend_kompute_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) {
     GGML_UNUSED(buffer);
 
-    const auto res = ggml_vk_get_tensor(tensor);
+    const auto res = ggml_vk_get_tensor_slice(tensor, offset, size);
     GGML_ASSERT(res);
 
     komputeManager()->sequence()->eval<kp::OpTensorSyncLocal>({res});

llama.cpp

@@ -16447,7 +16447,7 @@ static void llama_state_get_data_internal(struct llama_context * ctx, llama_data
         data_ctx->write(&kv_used,     sizeof(kv_used));
         data_ctx->write(&v_trans,     sizeof(v_trans));
 
-        if (kv_buf_size) {
+        if (kv_buf_size && kv_head) {
             const size_t pre_kv_buf_size = data_ctx->get_size_written();
 
             std::vector<uint8_t> tmp_buf;
@@ -16611,10 +16611,10 @@ size_t llama_state_set_data(struct llama_context * ctx, const uint8_t * src) {
 
         llama_kv_cache_clear(ctx);
 
-        if (kv_buf_size) {
-            const size_t pre_kv_buf_size = inp - src;
+        GGML_ASSERT(kv_self.total_size() >= kv_buf_size);
 
-            GGML_ASSERT(kv_self.total_size() >= kv_buf_size);
+        if (kv_buf_size && kv_head) {
+            const size_t pre_kv_buf_size = inp - src;
 
             for (int il = 0; il < (int) n_layer; ++il) {
                 const size_t k_size = ggml_row_size(kv_self.k_l[il]->type, n_embd_k_gqa*kv_head);