From 25f9b152f919b0262531873c935c301201d546aa Mon Sep 17 00:00:00 2001 From: Jesse Gross Date: Fri, 7 Mar 2025 17:20:54 -0800 Subject: [PATCH] ggml-backend: Ensure allocation meet backend requirements Backends can impose additional alignment requirements on buffer sizes. We should ensure that we meet these or allocations can fail. --- ml/backend/ggml/ggml.go | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/ml/backend/ggml/ggml.go b/ml/backend/ggml/ggml.go index 00873b4fc..bf17c8824 100644 --- a/ml/backend/ggml/ggml.go +++ b/ml/backend/ggml/ggml.go @@ -520,6 +520,10 @@ func shapeToGGML(shape []int) *C.int64_t { return &sh[0] } +func pad(length, pad C.size_t) C.size_t { + return ((length + pad - 1) / pad) * pad +} + func (c Context) newTensor(dtype ml.DType, shape []int) ml.Tensor { if c.buft == nil { panic("set Input, Output, or Layer before creating tensors") @@ -551,7 +555,8 @@ func (c Context) newTensor(dtype ml.DType, shape []int) ml.Tensor { } t := C.ggml_new_tensor(c.ctx, cdtype, C.int(len(shape)), shapeToGGML(shape)) - b := C.ggml_backend_buft_alloc_buffer(c.buft, C.ggml_nbytes(t)) + size := pad(C.ggml_backend_buft_get_alloc_size(c.buft, t), C.ggml_backend_buft_get_alignment(c.buft)) + b := C.ggml_backend_buft_alloc_buffer(c.buft, size) C.ggml_backend_tensor_alloc(b, t, C.ggml_backend_buffer_get_base(b)) return &Tensor{b: c.b, t: t} }