Merge pull request #1180 from dpmm99/feat/tensor-override

Feat/tensor override

Author: martindevans

.github/_typos.toml

@@ -17,4 +17,6 @@ extend-exclude = [
 
 [default.extend-words]
 # Used in a comment in SafeLLamaSamplerHandle.cs, as a prefix of "hello"
-teh = "hel"
+teh = "hel"
+# ot is the shorthand version of llama.cpp's override-tensor parameter
+ot = "ot"

LLama.Unittest/ModelsParamsTests.cs

@@ -41,6 +41,11 @@ public void SerializeRoundTripSystemTextJson()
             actual.MetadataOverrides = null!;
             expected.MetadataOverrides = null!;
 
+            // Same deal
+            Assert.True(expected.TensorBufferOverrides.SequenceEqual(actual.TensorBufferOverrides));
+            actual.TensorBufferOverrides = null!;
+            expected.TensorBufferOverrides = null!;
+
             // Check encoding is the same
             var b1 = expected.Encoding.GetBytes("Hello");
             var b2 = actual.Encoding.GetBytes("Hello");

LLama.Web/Common/ModelOptions.cs

@@ -26,6 +26,9 @@ public class ModelOptions
         /// <inheritdoc />
         public GPUSplitMode? SplitMode { get; set; }
 
+        /// <inheritdoc />
+        public List<TensorBufferOverride> TensorBufferOverrides { get; set; } = new();
+
         /// <inheritdoc />
         public int GpuLayerCount { get; set; } = 20;
 

LLama/Abstractions/IModelParams.cs

@@ -38,6 +38,12 @@ public interface IModelParams
         /// </summary>
         GPUSplitMode? SplitMode { get; }
 
+        /// <summary>
+        /// Buffer type overrides for specific tensor patterns, allowing you to specify hardware devices to use for individual tensors or sets of tensors.
+        /// Equivalent to --override-tensor or -ot on the llama.cpp command line or tensor_buft_overrides internally.
+        /// </summary>
+        List<TensorBufferOverride> TensorBufferOverrides { get; }
+
         /// <summary>
         /// Number of layers to run in VRAM / GPU memory (n_gpu_layers)
         /// </summary>

LLama/Abstractions/TensorBufferOverride.cs

@@ -0,0 +1,36 @@
+using System;
+
+namespace LLama.Abstractions
+{
+    /// <summary>
+    /// Represents a mapping between a tensor name pattern and a specific buffer type
+    /// </summary>
+    public class TensorBufferOverride
+    {
+        /// <summary>
+        /// Pattern to match tensor names. This is a regular expression. You can check the tensor names via the model.Metadata.
+        /// </summary>
+        public string Pattern { get; set; }
+
+        /// <summary>
+        /// Buffer type to use for matching tensors. Examples: CPU, GPU0, GPU1
+        /// </summary>
+        public string BufferType { get; set; }
+
+        /// <summary>
+        /// Creates a new tensor buffer override
+        /// </summary>
+        /// <param name="pattern">Pattern to match tensor names</param>
+        /// <param name="bufferType">Buffer type to use for matching tensors</param>
+        public TensorBufferOverride(string pattern, string bufferType)
+        {
+            if (string.IsNullOrEmpty(pattern))
+                throw new ArgumentException("Pattern cannot be null or empty", nameof(pattern));
+            if (string.IsNullOrEmpty(bufferType))
+                throw new ArgumentException("Buffer type cannot be null or empty", nameof(bufferType));
+
+            Pattern = pattern;
+            BufferType = bufferType;
+        }
+    }
+}

LLama/Common/ModelParams.cs

@@ -21,6 +21,9 @@ public record ModelParams
         /// <inheritdoc />
         public GPUSplitMode? SplitMode { get; set; }
 
+        /// <inheritdoc />
+        public List<TensorBufferOverride> TensorBufferOverrides { get; set; } = new();
+
         /// <inheritdoc />
         public int GpuLayerCount { get; set; } = 20;
 

LLama/Extensions/IModelParamsExtensions.cs

@@ -45,6 +45,20 @@ public static IDisposable ToLlamaModelParams(this IModelParams @params, out LLam
             result.tensor_split = (float*)disposer.Add(@params.TensorSplits.Pin()).Pointer;
         }
 
+        // Add tensor buffer overrides, if any
+        if (@params.TensorBufferOverrides.Count > 0)
+        {
+            var bufferOverrideHelper = new LLamaTensorBufferOverrideHelper();
+            disposer.Add(bufferOverrideHelper);
+
+            foreach (var tensorOverride in @params.TensorBufferOverrides)
+            {
+                bufferOverrideHelper.AddOverride(tensorOverride.Pattern, tensorOverride.BufferType);
+            }
+
+            bufferOverrideHelper.ApplyToModelParams(ref result);
+        }
+
         if (@params.MetadataOverrides.Count == 0)
         {
             unsafe

LLama/Native/LLamaModelParams.cs

@@ -13,11 +13,11 @@ public unsafe struct LLamaModelParams
         /// todo: add support for llama_model_params.devices
         /// </summary>
         private IntPtr devices;
-
-        // NULL-terminated list of buffer types to use for tensors that match a pattern
-        // actual type: llama_model_tensor_buft_override* 
-        // todo: add support for tensor_buft_overrides
-        private IntPtr tensor_buft_overrides;
+
+        /// <summary>
+        /// NULL-terminated list of buffer types to use for tensors that match a pattern
+        /// </summary>
+        public LLamaModelTensorBufferOverride* tensor_buft_overrides;
 
         /// <summary>
         /// // number of layers to store in VRAM

LLama/Native/LLamaModelTensorBufferOverride.cs

@@ -0,0 +1,22 @@
+using System;
+
+namespace LLama.Native
+{
+    /// <summary>
+    /// Represents a mapping between a tensor name pattern and a backend buffer type<br/>
+    /// Original type: llama_model_tensor_buft_override
+    /// </summary>
+    [StructLayout(LayoutKind.Sequential)]
+    public struct LLamaModelTensorBufferOverride
+    {
+        /// <summary>
+        /// Tensor name pattern to match
+        /// </summary>
+        public IntPtr Pattern;
+
+        /// <summary>
+        /// Backend buffer type to use for matching tensors, as obtained via ggml_backend_dev_buffer_type
+        /// </summary>
+        public IntPtr BufferType;
+    }
+}

LLama/Native/LLamaTensorBufferOverrideHelper.cs

@@ -0,0 +1,135 @@
+using System;
+using System.Collections.Generic;
+using System.Text;
+
+namespace LLama.Native
+{
+    /// <summary>
+    /// Helper for creating and managing tensor buffer overrides
+    /// </summary>
+    internal class LLamaTensorBufferOverrideHelper : IDisposable
+    {
+        private readonly List<IntPtr> _allocatedMemory = new();
+        private readonly List<LLamaModelTensorBufferOverride> _overrides = new();
+        private IntPtr _overrideArray = IntPtr.Zero;
+        private readonly Dictionary<string, IntPtr> _bufferTypeCache = new();
+
+        /// <summary>
+        /// Get all available buffer types
+        /// </summary>
+        /// <returns>Dictionary mapping buffer type names to their handles</returns>
+        public Dictionary<string, IntPtr> GetAvailableBufferTypes()
+        {
+            var result = new Dictionary<string, IntPtr>();
+            
+            nuint count = NativeApi.ggml_backend_dev_count();
+            for (nuint i = 0; i < count; i++)
+            {
+                IntPtr dev = NativeApi.ggml_backend_dev_get(i);
+                IntPtr buft = NativeApi.ggml_backend_dev_buffer_type(dev);
+                
+                if (buft != IntPtr.Zero)
+                {
+                    IntPtr namePtr = NativeApi.ggml_backend_buft_name(buft);
+                    string name = Marshal.PtrToStringAnsi(namePtr) ?? string.Empty;
+                    
+                    if (!string.IsNullOrEmpty(name))
+                    {
+                        result[name] = buft;
+                        _bufferTypeCache[name] = buft;
+                    }
+                }
+            }
+            
+            return result;
+        }
+
+        /// <summary>
+        /// Add a tensor buffer override
+        /// </summary>
+        /// <param name="pattern">Tensor name pattern to match</param>
+        /// <param name="bufferTypeName">Name of the buffer type to use</param>
+        /// <returns>True if the override was added successfully</returns>
+        public bool AddOverride(string pattern, string bufferTypeName)
+        {
+            if (string.IsNullOrEmpty(pattern) || string.IsNullOrEmpty(bufferTypeName))
+                return false;
+
+            // Get all buffer types if cache is empty
+            if (_bufferTypeCache.Count == 0)
+            {
+                GetAvailableBufferTypes();
+            }
+
+            // Check if we have this buffer type
+            if (!_bufferTypeCache.TryGetValue(bufferTypeName, out IntPtr bufferType))
+                return false;
+
+            // Allocate memory for the pattern string and keep track of it
+            byte[] patternBytes = Encoding.UTF8.GetBytes(pattern + "\0");
+            IntPtr patternPtr = Marshal.AllocHGlobal(patternBytes.Length);
+            Marshal.Copy(patternBytes, 0, patternPtr, patternBytes.Length);
+            _allocatedMemory.Add(patternPtr);
+
+            // Create the override
+            var @override = new LLamaModelTensorBufferOverride
+            {
+                Pattern = patternPtr,
+                BufferType = bufferType
+            };
+
+            _overrides.Add(@override);
+            return true;
+        }
+
+        /// <summary>
+        /// Apply the overrides to model parameters
+        /// </summary>
+        /// <param name="modelParams">Model parameters to update</param>
+        public unsafe void ApplyToModelParams(ref LLamaModelParams modelParams)
+        {
+            if (_overrides.Count == 0)
+            {
+                modelParams.tensor_buft_overrides = null;
+                return;
+            }
+
+            // Free previous array if it exists
+            if (_overrideArray != IntPtr.Zero)
+            {
+                Marshal.FreeHGlobal(_overrideArray);
+            }
+
+            // Allocate memory for the array + null terminator
+            int size = Marshal.SizeOf<LLamaModelTensorBufferOverride>() * (_overrides.Count + 1);
+            _overrideArray = Marshal.AllocHGlobal(size);
+            _allocatedMemory.Add(_overrideArray);
+
+            // Copy overrides to array
+            for (int i = 0; i < _overrides.Count; i++)
+            {
+                IntPtr elemPtr = IntPtr.Add(_overrideArray, i * Marshal.SizeOf<LLamaModelTensorBufferOverride>());
+                Marshal.StructureToPtr(_overrides[i], elemPtr, false);
+            }
+
+            // Add null terminator
+            IntPtr nullTermPtr = IntPtr.Add(_overrideArray, _overrides.Count * Marshal.SizeOf<LLamaModelTensorBufferOverride>());
+            Marshal.StructureToPtr(new LLamaModelTensorBufferOverride { Pattern = IntPtr.Zero, BufferType = IntPtr.Zero }, nullTermPtr, false);
+
+            // Update model params
+            modelParams.tensor_buft_overrides = (LLamaModelTensorBufferOverride*)_overrideArray.ToPointer();
+        }
+
+        /// <inheritdoc />
+        public void Dispose()
+        {
+            foreach (IntPtr ptr in _allocatedMemory)
+            {
+                Marshal.FreeHGlobal(ptr);
+            }
+            _allocatedMemory.Clear();
+            _overrides.Clear();
+            _overrideArray = IntPtr.Zero;
+        }
+    }
+}