- Added low level wrapper methods for new per-sequence state load/save in SafeLLamaContextHandle

 - Added high level wrapper methods (save/load with `State` object or memory mapped file) in `LLamaContext`
 - Moved native methods for per-sequence state load/save into `SafeLLamaContextHandle`

Author: martindevans

LLama/LLamaContext.cs

@@ -152,6 +152,7 @@ public string DeTokenize(IReadOnlyList<LLamaToken> tokens)
             return decoder.Read();
         }
 
+        #region state load/save
         /// <summary>
         /// Save the state to specified path.
         /// </summary>
@@ -163,7 +164,7 @@ public void SaveState(string filename)
                 File.Delete(filename);
 
             // Estimate size of state to write to disk, this is always equal to or greater than the actual size
-            var estimatedStateSize = (long)NativeApi.llama_state_get_size(NativeHandle);
+            var estimatedStateSize = checked((long)NativeHandle.GetStateSize());
 
             // Map the file and write the bytes directly to it. This saves copying the bytes into a C# array
             long writtenBytes;
@@ -174,8 +175,53 @@ public void SaveState(string filename)
                 {
                     byte* ptr = null;
                     view.SafeMemoryMappedViewHandle.AcquirePointer(ref ptr);
-                    writtenBytes = (long)NativeApi.llama_state_get_data(NativeHandle, ptr);
-                    view.SafeMemoryMappedViewHandle.ReleasePointer();
+                    try
+                    {
+                        writtenBytes = (long)NativeHandle.GetState(ptr, (ulong)estimatedStateSize);
+                    }
+                    finally
+                    {
+                        view.SafeMemoryMappedViewHandle.ReleasePointer();
+                    }
+                }
+            }
+
+            // Truncate the file to the actual size of data that was written
+            using (var fileStream = new FileStream(filename, FileMode.Open))
+                fileStream.SetLength(writtenBytes);
+        }
+
+        /// <summary>
+        /// Save the state of a particular sequence to specified path.
+        /// </summary>
+        /// <param name="filename"></param>
+        /// <param name="sequence"></param>
+        public void SaveState(string filename, LLamaSeqId sequence)
+        {
+            // Delete that file before overwriting it
+            if (File.Exists(filename))
+                File.Delete(filename);
+
+            // Estimate size of state to write to disk, this is always equal to or greater than the actual size
+            var estimatedStateSize = checked((long)NativeHandle.GetStateSize(sequence));
+
+            // Map the file and write the bytes directly to it. This saves copying the bytes into a C# array
+            long writtenBytes;
+            using (var file = MemoryMappedFile.CreateFromFile(filename, FileMode.Create, null, estimatedStateSize))
+            using (var view = file.CreateViewAccessor(0, estimatedStateSize))
+            {
+                unsafe
+                {
+                    byte* ptr = null;
+                    view.SafeMemoryMappedViewHandle.AcquirePointer(ref ptr);
+                    try
+                    {
+                        writtenBytes = (long)NativeHandle.GetState(ptr, (ulong)estimatedStateSize);
+                    }
+                    finally
+                    {
+                        view.SafeMemoryMappedViewHandle.ReleasePointer();
+                    }
                 }
             }
 
@@ -187,7 +233,7 @@ public void SaveState(string filename)
         /// <summary>
         /// Get the state data as an opaque handle, which can be loaded later using <see cref="LoadState(State)"/>
         /// </summary>
-        /// <remarks>Use <see cref="SaveState"/> if you intend to save this state to disk.</remarks>
+        /// <remarks>Use <see cref="SaveState(string)"/> if you intend to save this state to disk.</remarks>
         /// <returns></returns>
         public State GetState()
         {
@@ -198,7 +244,11 @@ public State GetState()
             try
             {
                 // Copy the state data into memory, discover the actual size required
-                var actualSize = NativeHandle.GetState(memory, stateSize);
+                ulong actualSize;
+                unsafe
+                {
+                    actualSize = NativeHandle.GetState((byte*)memory, stateSize);
+                }
 
                 // Shrink to size
                 memory = Marshal.ReAllocHGlobal(memory, (nint)actualSize);
@@ -218,11 +268,48 @@ public State GetState()
             }
         }
 
+        /// <summary>
+        /// Get the state data as an opaque handle, which can be loaded later using <see cref="LoadState(State)"/>
+        /// </summary>
+        /// <remarks>Use <see cref="SaveState(string, LLamaSeqId)"/> if you intend to save this state to disk.</remarks>
+        /// <returns></returns>
+        public SequenceState GetState(LLamaSeqId sequence)
+        {
+            var stateSize = NativeHandle.GetStateSize(sequence);
+
+            // Allocate a chunk of memory large enough to hold the entire state
+            var memory = Marshal.AllocHGlobal((nint)stateSize);
+            try
+            {
+                // Copy the state data into memory, discover the actual size required
+                ulong actualSize;
+                unsafe
+                {
+                    actualSize = NativeHandle.GetState((byte*)memory, stateSize, sequence);
+                }
+
+                // Shrink to size
+                memory = Marshal.ReAllocHGlobal(memory, (nint)actualSize);
+
+                // Wrap memory in a "state"
+                var state = new SequenceState(memory, actualSize);
+
+                // Set memory to zero, to prevent it being freed in finally block
+                memory = IntPtr.Zero;
+
+                return state;
+            }
+            finally
+            {
+                if (memory != IntPtr.Zero)
+                    Marshal.FreeHGlobal(memory);
+            }
+        }
+
         /// <summary>
         /// Load the state from specified path.
         /// </summary>
         /// <param name="filename"></param>
-        /// <exception cref="RuntimeError"></exception>
         public void LoadState(string filename)
         {
             // Map state file into memory and pass that pointer directly to `llama_set_state_data` to load from
@@ -233,8 +320,41 @@ public void LoadState(string filename)
                 {
                     byte* ptr = null;
                     view.SafeMemoryMappedViewHandle.AcquirePointer(ref ptr);
-                    NativeApi.llama_state_set_data(NativeHandle, ptr);
-                    view.SafeMemoryMappedViewHandle.ReleasePointer();
+                    try
+                    {
+                        NativeHandle.SetState(ptr);
+                    }
+                    finally
+                    {
+                        view.SafeMemoryMappedViewHandle.ReleasePointer();
+                    }
+                }
+            }
+        }
+
+        /// <summary>
+        /// Load the state from specified path into a particular sequence
+        /// </summary>
+        /// <param name="filename"></param>
+        /// <param name="sequence"></param>
+        public void LoadState(string filename, LLamaSeqId sequence)
+        {
+            // Map state file into memory and pass that pointer directly to `llama_set_state_data` to load from
+            using (var file = MemoryMappedFile.CreateFromFile(filename, FileMode.Open, null))
+            using (var view = file.CreateViewAccessor())
+            {
+                unsafe
+                {
+                    byte* ptr = null;
+                    view.SafeMemoryMappedViewHandle.AcquirePointer(ref ptr);
+                    try
+                    {
+                        NativeHandle.SetState(ptr, sequence);
+                    }
+                    finally
+                    {
+                        view.SafeMemoryMappedViewHandle.ReleasePointer();
+                    }
                 }
             }
         }
@@ -248,10 +368,25 @@ public void LoadState(State state)
         {
             unsafe
             {
-                NativeHandle.SetState((byte*)state.DangerousGetHandle().ToPointer());
+                NativeHandle.SetState((byte*)state.DangerousGetHandle());
             }
         }
 
+        /// <summary>
+        /// Load the state from memory into a particular sequence
+        /// </summary>
+        /// <param name="state"></param>
+        /// <param name="sequence"></param>
+        /// <exception cref="RuntimeError"></exception>
+        public void LoadState(SequenceState state, LLamaSeqId sequence)
+        {
+            unsafe
+            {
+                NativeHandle.SetState((byte*)state.DangerousGetHandle(), sequence);
+            }
+        }
+        #endregion
+
         /// <summary>
         /// Sample a single token from this context, using the given sampling pipeline
         /// </summary>
@@ -417,12 +552,16 @@ public void Dispose()
         }
 
         /// <summary>
-        /// The state of this model, which can be reloaded later
+        /// The state of this context, which can be reloaded later
         /// </summary>
         public class State
             : SafeLLamaHandleBase
         {
-            private ulong _size;
+            private readonly ulong _size;
+            /// <summary>
+            /// Get the size in bytes of this state object
+            /// </summary>
+            public ulong Size => _size;
 
             internal State(IntPtr memory, ulong size)
                 : base(memory, true)
@@ -441,6 +580,7 @@ protected override bool ReleaseHandle()
             /// Convert this state to a byte array
             /// </summary>
             /// <returns></returns>
+            [Obsolete("It is not generally safe to convert a state into a byte array - it will fail if the state is very large")]
             public byte[] ToByteArray()
             {
                 var bytes = new byte[_size];
@@ -453,12 +593,57 @@ public byte[] ToByteArray()
             /// </summary>
             /// <param name="bytes"></param>
             /// <returns></returns>
+            [Obsolete("It is not generally safe to convert a state into a byte array - it will fail if the state is very large")]
             public static State FromByteArray(byte[] bytes)
             {
                 var memory = Marshal.AllocHGlobal(bytes.Length);
                 Marshal.Copy(bytes, 0, memory, bytes.Length);
                 return new State(memory, (ulong)bytes.Length);
             }
         }
+
+        /// <summary>
+        /// The state of a single sequence, which can be reloaded later
+        /// </summary>
+        public class SequenceState
+            : SafeLLamaHandleBase
+        {
+            private readonly ulong _size;
+            /// <summary>
+            /// Get the size in bytes of this state object
+            /// </summary>
+            public ulong Size => _size;
+
+            internal SequenceState(IntPtr memory, ulong size)
+                : base(memory, true)
+            {
+                _size = size;
+            }
+
+            /// <inheritdoc />
+            protected override bool ReleaseHandle()
+            {
+                Marshal.FreeHGlobal(handle);
+                return true;
+            }
+
+            /// <summary>
+            /// Copy bytes to a desintation pointer.
+            /// </summary>
+            /// <param name="dst">Destination to write to</param>
+            /// <param name="length">Length of the destination buffer</param>
+            /// <param name="offset">Offset from start of src to start copying from</param>
+            /// <returns>Number of bytes written to destination</returns>
+            public unsafe ulong CopyTo(byte* dst, ulong length, ulong offset = 0)
+            {
+                var copy = Math.Min(length, _size - offset);
+
+                var src = (byte*)DangerousGetHandle();
+                src += offset;
+
+                Buffer.MemoryCopy(src, dst, length, copy);
+                return copy;
+            }
+        }
     }
 }

LLama/Native/NativeApi.cs

@@ -67,34 +67,6 @@ public static void llama_empty_call()
         //[DllImport(libraryName, CallingConvention = CallingConvention.Cdecl)]
         //public static extern void llama_numa_init(ggml_numa_strategy numa);
 
-        /// <summary>
-        /// Returns the maximum size in bytes of the state (rng, logits, embedding
-        /// and kv_cache) - will often be smaller after compacting tokens
-        /// </summary>
-        /// <param name="ctx"></param>
-        /// <returns></returns>
-        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl)]
-        public static extern ulong llama_state_get_size(SafeLLamaContextHandle ctx);
-
-        /// <summary>
-        /// Copies the state to the specified destination address.
-        /// Destination needs to have allocated enough memory.
-        /// </summary>
-        /// <param name="ctx"></param>
-        /// <param name="dest"></param>
-        /// <returns>the number of bytes copied</returns>
-        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl)]
-        public static extern unsafe ulong llama_state_get_data(SafeLLamaContextHandle ctx, byte* dest);
-
-        /// <summary>
-        /// Set the state reading from the specified address
-        /// </summary>
-        /// <param name="ctx"></param>
-        /// <param name="src"></param>
-        /// <returns>the number of bytes read</returns>
-        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl)]
-        public static extern unsafe ulong llama_state_set_data(SafeLLamaContextHandle ctx, byte* src);
-
         /// <summary>
         /// Load session file
         /// </summary>
@@ -118,35 +90,6 @@ public static void llama_empty_call()
         [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl)]
         public static extern bool llama_state_save_file(SafeLLamaContextHandle ctx, string path_session, LLamaToken[] tokens, ulong n_token_count);
 
-        /// <summary>
-        /// Get the exact size needed to copy the KV cache of a single sequence
-        /// </summary>
-        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl)]
-        public static extern nuint llama_state_seq_get_size(SafeLLamaContextHandle ctx, LLamaSeqId seq_id);
-
-        /// <summary>
-        /// Copy the KV cache of a single sequence into the specified buffer
-        /// </summary>
-        /// <param name="ctx"></param>
-        /// <param name="dst"></param>
-        /// <param name="seq_id"></param>
-        /// <returns></returns>
-        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl)]
-        public static extern unsafe nuint llama_state_seq_get_data(SafeLLamaContextHandle ctx, byte* dst, LLamaSeqId seq_id);
-
-        /// <summary>
-        /// Copy the sequence data (originally copied with `llama_state_seq_get_data`) into the specified sequence
-        /// </summary>
-        /// <param name="ctx"></param>
-        /// <param name="src"></param>
-        /// <param name="dest_seq_id"></param>
-        /// <returns>
-        ///  - Positive: Ok
-        ///  - Zero: Failed to load
-        /// </returns>
-        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl)]
-        public static extern unsafe nuint llama_state_seq_set_data(SafeLLamaContextHandle ctx, byte* src, LLamaSeqId dest_seq_id);
-
         [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl)]
         public static extern unsafe nuint llama_state_seq_save_file(SafeLLamaContextHandle ctx, string filepath, LLamaSeqId seq_id, LLamaToken* tokens, nuint n_token_count);