Repository: ggml-org/llama.cpp
Pull request: ggml-alloc : fix reuse-parent logic for misaligned sizes

cont #16679

I encountered a bug in the logic for parent buffer reuse which occurs when the alloc size of the node and the parent are not aligned to the buffer alignment.

The logic in `ggml_gallocr_free_extra_space` aims to free the extra space that is left when reusing a parent buffer with a larger alloc size than the node. This case occurs currently only with the Metal backend because we allocate additional size for some tensors:

https://github.com/ggml-org/llama.cpp/blob/0cdce38a97bd05462416272fbb912f4d7ecd2940/ggml/src/ggml-metal/ggml-metal.cpp#L183-L217

The implementation of `ggml_gallocr_free_extra_space` did not take into account that the passed `extra_size` to the `ggml_dyn_tallocr_free_tensor` call will be padded/aligned here:

https://github.com/ggml-org/llama.cpp/blob/0cdce38a97bd05462416272fbb912f4d7ecd2940/ggml/src/ggml-alloc.c#L312-L316

Depending on the specific tensor sizes / alignment, this could lead to corruption of the allocator chunks. The fix in this PR takes into account the alignment before computing the `extra_size`, which guarantees that after freeing it, the chunks will remain aligned.

Also:

- Rename `ggml_dyn_tallocr_free_tensor` -> `ggml_dyn_tallocr_free_bytes`
- Fix crash when building with `GGML_ALLOCATOR_DEBUG` and using the Metal backend. The problem was that `ggml_dyn_tallocr_free_tensor` assumed to always remove an allocated tensor and hence the `remove_allocated_tensor()` debug call in it. However, when we free extra bytes from a parent buffer, we are not actually removing a tensor - just the extra bytes. To fix that, we now call `remove_allocated_tensor()` only in `ggml_gallocr_free_node()`

Changed files:
filename: ggml/src/ggml-alloc.c
patch: @@ -312,16 +312,9 @@ static struct buffer_address ggml_dyn_tallocr_alloc(struct ggml_dyn_tallocr * al
 }
 
 // this is a very naive implementation, but for our case the number of free blocks should be very small
-static void ggml_dyn_tallocr_free_tensor(struct ggml_dyn_tallocr * alloc, struct buffer_address addr, size_t size, const struct ggml_tensor * tensor) {
+static void ggml_dyn_tallocr_free_bytes(struct ggml_dyn_tallocr * alloc, struct buffer_address addr, size_t size) {
     size = aligned_offset(NULL, size, alloc->alignment);
 
-    AT_PRINTF("%s: freeing %s at {chunk=%d, offset=%zu} (%zu bytes) - n_free_blocks = %d\n",
-        __func__, tensor->name, addr.chunk, addr.offset, size, alloc->chunks[addr.chunk]->n_free_blocks);
-
-#ifdef GGML_ALLOCATOR_DEBUG
-    remove_allocated_tensor(alloc, addr, tensor);
-#endif
-
     struct tallocr_chunk * chunk = alloc->chunks[addr.chunk];
 
     // see if we can merge with an existing block
@@ -357,8 +350,6 @@ static void ggml_dyn_tallocr_free_tensor(struct ggml_dyn_tallocr * alloc, struct
     }
     // otherwise, add a new block
     ggml_dyn_tallocr_insert_block(chunk, addr.offset, size);
-
-    GGML_UNUSED(tensor);
 }
 
 static void ggml_dyn_tallocr_reset(struct ggml_dyn_tallocr * alloc) {
@@ -616,13 +607,17 @@ static void ggml_gallocr_free_extra_space(ggml_gallocr_t galloc, struct ggml_ten
 
     GGML_ASSERT(parent_size >= node_size);
 
+    // note: we want after the freeing the chunks to continue to be aligned
+    struct ggml_dyn_tallocr * p_alloc = galloc->buf_tallocs[p_hn->buffer_id];
+    parent_size = aligned_offset(NULL, parent_size, p_alloc->alignment);
+    node_size = aligned_offset(NULL, node_size, p_alloc->alignment);
+
     if (parent_size > node_size) {
-        struct ggml_dyn_tallocr * p_alloc = galloc->buf_tallocs[p_hn->buffer_id];
         struct buffer_address p_addr = p_hn->addr;
         p_addr.offset += node_size;
         size_t extra_size = parent_size - node_size;
         AT_PRINTF("freeing extra %zu bytes from parent %s for %s\n", extra_size, parent->name, node->name);
-        ggml_dyn_tallocr_free_tensor(p_alloc, p_addr, extra_size, parent);
+        ggml_dyn_tallocr_free_bytes(p_alloc, p_addr, extra_size);
     }
 }
 
@@ -706,7 +701,14 @@ static void ggml_gallocr_free_node(ggml_gallocr_t galloc, struct ggml_tensor * n
     struct ggml_dyn_tallocr * alloc = galloc->buf_tallocs[buffer_id];
     ggml_backend_buffer_type_t buft = galloc->bufts[buffer_id];
     size_t size = ggml_backend_buft_get_alloc_size(buft, node);
-    ggml_dyn_tallocr_free_tensor(alloc, hn->addr, size, node);
+
+    AT_PRINTF("%s: freeing %s at {chunk=%d, offset=%zu} (%zu bytes) - n_free_blocks = %d\n",
+        __func__, node->name, hn->addr.chunk, hn->addr.offset, size, alloc->chunks[hn->addr.chunk]->n_free_blocks);
+#ifdef GGML_ALLOCATOR_DEBUG
+    remove_allocated_tensor(alloc, hn->addr, node);
+#endif
+
+    ggml_dyn_tallocr_free_bytes(alloc, hn->addr, size);
     hn->allocated = false;
 }