From b0acecf080764684040a916ce7a6d7c79876ab45 Mon Sep 17 00:00:00 2001 From: Martin Evans Date: Sun, 21 Jan 2024 00:54:07 +0000 Subject: [PATCH] Created a new `BatchedExecutor` which processes multiple "Conversations" in one single inference batch. This is faster, even when the conversations are unrelated, and is much faster if the conversations share some overlap (e.g. a common system prompt prefix). Conversations can be "forked", to create a copy of a conversation at a given point. This allows e.g. prompting a conversation with a system prefix just once and then forking it again and again for each individual conversation. Conversations can also be "rewound" to an earlier state. Added two new examples, demonstrating forking and rewinding. --- LLama.Examples/Examples/BatchedDecoding.cs | 172 --------------- .../Examples/BatchedExecutorFork.cs | 138 ++++++++++++ .../Examples/BatchedExecutorRewind.cs | 121 ++++++++++ LLama.Examples/Examples/Runner.cs | 3 +- LLama/Batched/BatchedExecutor.cs | 107 +++++++++ LLama/Batched/Conversation.cs | 206 ++++++++++++++++++ LLama/Batched/Exceptions.cs | 81 +++++++ LLama/LLamaContext.cs | 4 +- LLama/LLamaInstructExecutor.cs | 1 + LLama/LLamaInteractExecutor.cs | 1 + LLama/Native/LLamaBatch.cs | 55 ++++- LLama/Sampling/BaseSamplingPipeline.cs | 19 +- LLama/Sampling/DefaultSamplingPipeline.cs | 26 ++- LLama/Sampling/ISamplingPipeline.cs | 13 ++ 14 files changed, 748 insertions(+), 199 deletions(-) delete mode 100644 LLama.Examples/Examples/BatchedDecoding.cs create mode 100644 LLama.Examples/Examples/BatchedExecutorFork.cs create mode 100644 LLama.Examples/Examples/BatchedExecutorRewind.cs create mode 100644 LLama/Batched/BatchedExecutor.cs create mode 100644 LLama/Batched/Conversation.cs create mode 100644 LLama/Batched/Exceptions.cs diff --git a/LLama.Examples/Examples/BatchedDecoding.cs b/LLama.Examples/Examples/BatchedDecoding.cs deleted file mode 100644 index f893b613..00000000 --- a/LLama.Examples/Examples/BatchedDecoding.cs +++ /dev/null @@ -1,172 +0,0 @@ -using System.Diagnostics; -using System.Text; -using LLama.Common; -using LLama.Native; -using LLama.Sampling; - -namespace LLama.Examples.Examples; - -/// -/// This demonstrates generating multiple replies to the same prompt, with a shared cache -/// -/// Note that this is currently using the low level API directly, future work will provide a safer C# wrapper over this! -public class BatchedDecoding -{ - private const int n_parallel = 8; - private const int n_len = 32; - - public static async Task Run() - { - Console.Write("Please input your model path: "); - var modelPath = Console.ReadLine(); - - Console.WriteLine("Prompt (leave blank to select automatically):"); - var prompt = Console.ReadLine(); - if (string.IsNullOrWhiteSpace(prompt)) - prompt = "Not many people know that"; - - // Load model - var parameters = new ModelParams(modelPath); - - using var model = LLamaWeights.LoadFromFile(parameters); - - // Tokenize prompt - var prompt_tokens = model.Tokenize(prompt, true, false, Encoding.UTF8); - var n_kv_req = prompt_tokens.Length + (n_len - prompt_tokens.Length) * n_parallel; - - // Create a context - parameters.ContextSize = (uint)model.ContextSize; - parameters.BatchSize = (uint)Math.Max(n_len, n_parallel); - using var context = model.CreateContext(parameters); - - var n_ctx = context.ContextSize; - - // make sure the KV cache is big enough to hold all the prompt and generated tokens - if (n_kv_req > n_ctx) - { - await Console.Error.WriteLineAsync($"error: n_kv_req ({n_kv_req}) > n_ctx, the required KV cache size is not big enough\n"); - await Console.Error.WriteLineAsync(" either reduce n_parallel or increase n_ctx\n"); - return; - } - - var batch = new LLamaBatch(); - - // evaluate the initial prompt - batch.AddRange(prompt_tokens, 0, LLamaSeqId.Zero, true); - - if (await context.DecodeAsync(batch) != DecodeResult.Ok) - { - await Console.Error.WriteLineAsync("llama_decode failed"); - return; - } - - // assign the system KV cache to all parallel sequences - // this way, the parallel sequences will "reuse" the prompt tokens without having to copy them - for (var i = 1; i < n_parallel; ++i) - { - context.NativeHandle.KvCacheSequenceCopy((LLamaSeqId)0, (LLamaSeqId)i, 0, batch.TokenCount); - } - - if (n_parallel > 1) - { - Console.WriteLine(); - Console.WriteLine($"generating {n_parallel} sequences..."); - } - - // remember the batch index of the last token for each parallel sequence - // we need this to determine which logits to sample from - List i_batch = new(); - for (var i = 0; i < n_parallel; i++) - i_batch.Add(batch.TokenCount - 1); - - // Create per-stream decoder and sampler - var decoders = new StreamingTokenDecoder[n_parallel]; - var samplers = new ISamplingPipeline[n_parallel]; - for (var i = 0; i < n_parallel; i++) - { - decoders[i] = new StreamingTokenDecoder(context); - samplers[i] = new DefaultSamplingPipeline - { - Temperature = 0.1f + (float)i / n_parallel, - MinP = 0.25f, - }; - } - - var n_cur = batch.TokenCount; - var n_decode = 0; - - var timer = new Stopwatch(); - timer.Start(); - while (n_cur <= n_len) - { - batch.Clear(); - - for (var i = 0; i < n_parallel; i++) - { - // Skip completed streams - if (i_batch[i] < 0) - continue; - - // Use the sampling pipeline to select a token - var new_token_id = samplers[i].Sample( - context.NativeHandle, - context.NativeHandle.GetLogitsIth(i_batch[i]), - Array.Empty() - ); - - // Finish this stream early if necessary - if (new_token_id == model.EndOfSentenceToken || new_token_id == model.NewlineToken) - { - i_batch[i] = -1; - Console.WriteLine($"Completed Stream {i} early"); - continue; - } - - // Add this token to the decoder, so it will be turned into text - decoders[i].Add(new_token_id); - - i_batch[i] = batch.TokenCount; - - // push this new token for next evaluation - batch.Add(new_token_id, n_cur, (LLamaSeqId)i, true); - - n_decode++; - } - - // Check if all streams are finished - if (batch.TokenCount == 0) - { - break; - } - - n_cur++; - - // evaluate the current batch with the transformer model - if (await context.DecodeAsync(batch) != 0) - { - await Console.Error.WriteLineAsync("failed to eval"); - return; - } - } - - timer.Stop(); - Console.ForegroundColor = ConsoleColor.Yellow; - Console.WriteLine(); - Console.WriteLine($"Decoded {n_decode} tokens in {timer.ElapsedMilliseconds}ms"); - Console.WriteLine($"Rate: {n_decode / timer.Elapsed.TotalSeconds:##.000} tokens/second"); - - var index = 0; - foreach (var stream in decoders) - { - var text = stream.Read(); - - Console.ForegroundColor = ConsoleColor.Green; - Console.Write($"{index++}. {prompt}"); - Console.ForegroundColor = ConsoleColor.Red; - Console.WriteLine(text); - } - - Console.WriteLine("Press any key to exit demo"); - Console.ReadKey(true); - } -} \ No newline at end of file diff --git a/LLama.Examples/Examples/BatchedExecutorFork.cs b/LLama.Examples/Examples/BatchedExecutorFork.cs new file mode 100644 index 00000000..b834190c --- /dev/null +++ b/LLama.Examples/Examples/BatchedExecutorFork.cs @@ -0,0 +1,138 @@ +using LLama.Batched; +using LLama.Common; +using LLama.Native; +using LLama.Sampling; + +namespace LLama.Examples.Examples; + +/// +/// This demonstrates generating multiple replies to the same prompt, with a shared cache +/// +public class BatchedExecutorFork +{ + private const int n_split = 16; + private const int n_len = 64; + + public static async Task Run() + { + Console.Write("Please input your model path: "); + var modelPath = Console.ReadLine(); + + var parameters = new ModelParams(modelPath); + using var model = LLamaWeights.LoadFromFile(parameters); + + Console.WriteLine("Prompt (leave blank to select automatically):"); + var prompt = Console.ReadLine(); + if (string.IsNullOrWhiteSpace(prompt)) + prompt = "Not many people know that"; + + // Create an executor that can evaluate a batch of conversations together + var executor = new BatchedExecutor(model, parameters); + + // Print some info + var name = executor.Model.Metadata.GetValueOrDefault("general.name", "unknown model name"); + Console.WriteLine($"Created executor with model: {name}"); + + // Evaluate the initial prompt to create one conversation + var start = executor.Prompt(prompt); + await executor.Infer(); + + // Create the root node of the tree + var root = new Node(start); + + // Run inference loop + for (var i = 0; i < n_len; i++) + { + if (i != 0) + await executor.Infer(); + + // Occasionally fork all the active conversations + if (i != 0 && i % n_split == 0) + root.Split(); + + // Sample all active conversations + root.Sample(); + } + + Console.WriteLine($"{prompt}..."); + root.Print(1); + + Console.WriteLine("Press any key to exit demo"); + Console.ReadKey(true); + } + + class Node + { + private readonly StreamingTokenDecoder _decoder; + + private readonly DefaultSamplingPipeline _sampler; + private Conversation? _conversation; + + private Node? _left; + private Node? _right; + + public int ActiveConversationCount => _conversation != null ? 1 : _left!.ActiveConversationCount + _right!.ActiveConversationCount; + + public Node(Conversation conversation) + { + _sampler = new DefaultSamplingPipeline(); + _conversation = conversation; + _decoder = new StreamingTokenDecoder(conversation.Executor.Context); + } + + public void Sample() + { + if (_conversation == null) + { + _left?.Sample(); + _right?.Sample(); + return; + } + + if (_conversation.RequiresInference) + return; + + // Sample one token + var ctx = _conversation.Executor.Context.NativeHandle; + var logitsCopy = _conversation.Sample().ToArray(); + var token = _sampler.Sample(ctx, logitsCopy, Array.Empty()); + _sampler.Accept(ctx, token); + _decoder.Add(token); + + // Prompt the conversation with this token, to continue generating from there + _conversation.Prompt(token); + } + + public void Split() + { + if (_conversation != null) + { + _left = new Node(_conversation.Fork()); + _right = new Node(_conversation.Fork()); + + _conversation.Dispose(); + _conversation = null; + } + else + { + _left?.Split(); + _right?.Split(); + } + } + + public void Print(int indendation) + { + var colors = new[] { ConsoleColor.Red, ConsoleColor.Green, ConsoleColor.Blue, ConsoleColor.Yellow, ConsoleColor.White }; + Console.ForegroundColor = colors[indendation % colors.Length]; + + var message = _decoder.Read().ReplaceLineEndings(""); + + var prefix = new string(' ', indendation * 3); + var suffix = _conversation == null ? "..." : ""; + Console.WriteLine($"{prefix}...{message}{suffix}"); + + _left?.Print(indendation + 2); + _right?.Print(indendation + 2); + } + } +} \ No newline at end of file diff --git a/LLama.Examples/Examples/BatchedExecutorRewind.cs b/LLama.Examples/Examples/BatchedExecutorRewind.cs new file mode 100644 index 00000000..25195a56 --- /dev/null +++ b/LLama.Examples/Examples/BatchedExecutorRewind.cs @@ -0,0 +1,121 @@ +using LLama.Batched; +using LLama.Common; +using LLama.Native; +using LLama.Sampling; + +namespace LLama.Examples.Examples; + +/// +/// This demonstrates generating tokens and then rewinding to an earlier state +/// +public class BatchedExecutorRewind +{ + private const int n_generate = 24; + private const int n_rewind = 12; + private const int n_repeats = 6; + + public static async Task Run() + { + Console.Write("Please input your model path: "); + var modelPath = Console.ReadLine(); + + var parameters = new ModelParams(modelPath); + using var model = LLamaWeights.LoadFromFile(parameters); + + Console.WriteLine("Prompt (leave blank to select automatically):"); + var prompt = Console.ReadLine(); + if (string.IsNullOrWhiteSpace(prompt)) + prompt = "Not many people know that"; + + // Create an executor that can evaluate a batch of conversations together + var executor = new BatchedExecutor(model, parameters); + + // Print some info + var name = executor.Model.Metadata.GetValueOrDefault("general.name", "unknown model name"); + Console.WriteLine($"Created executor with model: {name}"); + + // Evaluate the initial prompt to create one conversation + var conversation = executor.Prompt(prompt); + + // Create the start node wrapping the conversation + var node = new Node(executor.Context); + + // Print the prompt + Console.ForegroundColor = ConsoleColor.Green; + Console.WriteLine(prompt); + + for (var i = 0; i < n_repeats; i++) + { + for (var j = 0; j < n_generate; j++) + { + // Run inference + await executor.Infer(); + + // Sample a token + var token = node.Sample(conversation); + + // Continue conversation with this token + if (j != n_generate - 1) + conversation.Prompt(token); + } + + // Write out what we generated + node.Write(n_rewind, i + 1); + + // Rewind back a few tokens + conversation.Rewind(n_rewind + 1); + + // Prompt with a token + conversation.Prompt(node.GetToken(n_generate - n_rewind - 1)); + + // Create a new node around the rewound conversation + node = new Node(executor.Context); + } + + Console.WriteLine("Press any key to exit demo"); + Console.ReadKey(true); + } + + private class Node + { + private readonly LLamaContext _context; + + private readonly List _tokens = new List(); + private readonly DefaultSamplingPipeline Sampler; + + public Node(LLamaContext context) + { + _context = context; + Sampler = new DefaultSamplingPipeline(); + } + + public LLamaToken Sample(Conversation conversation) + { + var token = Sampler.Sample(_context.NativeHandle, conversation.Sample().ToArray(), Array.Empty()); + _tokens.Add(token); + return token; + } + + public void Write(int n_rewind, int depth) + { + var decoder = new StreamingTokenDecoder(_context); + + for (var i = 0; i < _tokens.Count - n_rewind; i++) + decoder.Add(_tokens[i]); + + Console.ForegroundColor = ConsoleColor.Green; + Console.Write(new string(' ', depth * 3) + decoder.Read().ReplaceLineEndings(" ")); + + for (var i = _tokens.Count - n_rewind; i < _tokens.Count; i++) + decoder.Add(_tokens[i]); + + Console.ForegroundColor = ConsoleColor.DarkRed; + Console.WriteLine(decoder.Read().ReplaceLineEndings(" ")); + } + + public LLamaToken GetToken(int index) + { + return _tokens[index]; + } + } +} \ No newline at end of file diff --git a/LLama.Examples/Examples/Runner.cs b/LLama.Examples/Examples/Runner.cs index 3d9858e1..54358ce7 100644 --- a/LLama.Examples/Examples/Runner.cs +++ b/LLama.Examples/Examples/Runner.cs @@ -23,7 +23,8 @@ public class Runner { "Semantic Kernel Chat.", SemanticKernelChat.Run }, { "Semantic Kernel Memory.", SemanticKernelMemory.Run }, { "Coding Assistant.", CodingAssistant.Run }, - { "Batch Decoding.", BatchedDecoding.Run }, + { "Batched Executor (Fork)", BatchedExecutorFork.Run }, + { "Batched Executor (Rewind)", BatchedExecutorRewind.Run }, { "SK Kernel Memory.", KernelMemory.Run }, { "Exit", async () => Environment.Exit(0) } }; diff --git a/LLama/Batched/BatchedExecutor.cs b/LLama/Batched/BatchedExecutor.cs new file mode 100644 index 00000000..2fc3eabf --- /dev/null +++ b/LLama/Batched/BatchedExecutor.cs @@ -0,0 +1,107 @@ +using System; +using System.Threading; +using System.Threading.Tasks; +using LLama.Abstractions; +using LLama.Native; + +namespace LLama.Batched; + +/// +/// A batched executor that can infer multiple separate "conversations" simultaneously. +/// +public sealed class BatchedExecutor + : IDisposable +{ + private int _nextSequenceId; + + internal LLamaBatch Batch { get; } + internal ulong Epoch { get; private set; } + + /// + /// The this executor is using + /// + public LLamaContext Context { get; } + + /// + /// The this executor is using + /// + public LLamaWeights Model { get; } + + /// + /// Get the number of tokens in the batch, waiting for to be called + /// + public int BatchedTokenCount => Batch.TokenCount; + + /// + /// Check if this executor has been disposed. + /// + public bool IsDisposed { get; private set; } + + /// + /// Create a new batched executor + /// + /// The model to use + /// Parameters to create a new context + public BatchedExecutor(LLamaWeights model, IContextParams contextParams) + { + Model = model; + + Batch = new LLamaBatch(); + Context = model.CreateContext(contextParams); + } + + /// + /// Start a new with the given prompt + /// + /// + /// + public Conversation Prompt(string prompt) + { + if (IsDisposed) + throw new ObjectDisposedException(nameof(BatchedExecutor)); + + var conversation = new Conversation(this, GetNextSequenceId(), 0); + conversation.Prompt(prompt); + + return conversation; + } + + /// + /// Run inference for all conversations in the batch which have pending tokens. + /// + /// If the result is `NoKvSlot` then there is not enough memory for inference, try disposing some conversation + /// threads and running inference again. + /// + public async Task Infer(CancellationToken cancellation = default) + { + if (IsDisposed) + throw new ObjectDisposedException(nameof(BatchedExecutor)); + + var status = await Context.DecodeAsync(Batch, cancellation); + + // Only clear the batch if the result was ok. leaving all this state in place means that "Infer" can + // be called again after a warning (e.g. NoKvSlot). + if (status == DecodeResult.Ok) + { + Epoch++; + Batch.Clear(); + } + + return status; + } + + /// + public void Dispose() + { + if (IsDisposed) + return; + IsDisposed = true; + + Context.Dispose(); + } + + internal LLamaSeqId GetNextSequenceId() + { + return checked((LLamaSeqId)_nextSequenceId++); + } +} \ No newline at end of file diff --git a/LLama/Batched/Conversation.cs b/LLama/Batched/Conversation.cs new file mode 100644 index 00000000..3dcdec54 --- /dev/null +++ b/LLama/Batched/Conversation.cs @@ -0,0 +1,206 @@ +using System; +using System.Collections.Generic; +using LLama.Native; + +namespace LLama.Batched; + +/// +/// A single conversation thread that can be prompted (adding tokens from the user) or inferred (extracting a token from the LLM) +/// +public sealed class Conversation + : IDisposable +{ + private ulong _requiredEpoch; + private LLamaPos _end; + private int _batchIndex; + private bool _disposed; + + /// + /// The executor which this conversation belongs to + /// + public BatchedExecutor Executor { get; } + + /// + /// Unique ID for this conversation + /// + public LLamaSeqId ConversationId { get; } + + /// + /// Total number of tokens in this conversation, cannot exceed the context length. + /// + public int TokenCount => _end.Value; + + /// + /// Indicates if this conversation has been disposed, nothing can be done with a disposed conversation + /// + public bool IsDisposed => _disposed || Executor.IsDisposed; + + /// + /// Indicates if this conversation is waiting for inference to be run on the executor. "Prompt" and "Sample" cannot be called when this is true. + /// + public bool RequiresInference => _requiredEpoch > Executor.Epoch; + + /// + /// Indicates that this conversation should be sampled. + /// + public bool RequiresSampling => _requiredEpoch == Executor.Epoch; + + #region construction/destruction + internal Conversation(BatchedExecutor batch, LLamaSeqId id, LLamaPos end) + { + ConversationId = id; + Executor = batch; + + _end = end; + } + + /// + /// End this conversation, freeing all resources used by it + /// + /// + public void Dispose() + { + if (IsDisposed) + return; + _disposed = true; + + // Remove this conversation from the KV cache + Executor.Context.NativeHandle.KvCacheRemove(ConversationId, 0, _end); + } + + private void AssertNotDisposed() + { + if (Executor.IsDisposed) + throw new ObjectDisposedException(nameof(BatchedExecutor)); + if (IsDisposed) + throw new ObjectDisposedException(nameof(Conversation)); + } + #endregion + + /// + /// Create a copy of the current conversation + /// + /// The copy shares internal state, so consumes very little extra memory. + /// + /// + public Conversation Fork() + { + AssertNotDisposed(); + + if (RequiresInference) + throw new CannotForkWhileRequiresInference(); + + // Assign tokens to the new sequence + var id2 = Executor.GetNextSequenceId(); + NativeApi.llama_kv_cache_seq_cp(Executor.Context.NativeHandle, ConversationId, id2, 0, _end); + + // Create a new conversation which references the current position in this one + var c = new Conversation(Executor, id2, _end) + { + _batchIndex = _batchIndex, + _requiredEpoch = _requiredEpoch, + }; + + return c; + } + + /// + /// Rewind this conversation back to an earlier state + /// + /// + /// + /// + /// Thrown if `tokens` parameter is larger than NTokens + public void Rewind(int tokens) + { + AssertNotDisposed(); + + if (tokens > TokenCount) + throw new ArgumentOutOfRangeException(nameof(tokens), "Cannot rewind more than the total number of tokens"); + + // Remove those tokens from KV + Executor.Context.NativeHandle.KvCacheRemove(ConversationId, _end.Value - tokens, _end); + + // Adjust "end" marker back + _end = _end.Value - tokens; + } + + #region sample + /// + /// Get the logits from this conversation, ready for sampling + /// + /// + /// + /// Thrown if this conversation was not prompted before the previous call to infer + /// Thrown if Infer() must be called on the executor + public ReadOnlySpan Sample() + { + AssertNotDisposed(); + + if (_requiredEpoch < Executor.Epoch) + throw new CannotSampleRequiresPromptException(); + if (_requiredEpoch > Executor.Epoch) + throw new CannotSampleRequiresInferenceException(); + + return Executor.Context.NativeHandle.GetLogitsIth(_batchIndex); + } + #endregion + + #region prompt + private void AssertCanBePrompted() + { + AssertNotDisposed(); + + if (RequiresInference) + throw new AlreadyPromptedConversationException(); + } + + /// + /// Add tokens to this conversation + /// + /// + /// + public void Prompt(string input) + { + AssertCanBePrompted(); + + Prompt(Executor.Context.Tokenize(input)); + } + + /// + /// Add tokens to this conversation + /// + /// + /// + /// + public void Prompt(IReadOnlyList tokens) + { + AssertCanBePrompted(); + + // Add the prompt to the batch + for (var i = 0; i < tokens.Count; i++) + _batchIndex = Executor.Batch.Add(tokens[i], _end++, ConversationId, i == tokens.Count - 1); + + // Mark this conversation as needing inference/sampling + _requiredEpoch = Executor.Epoch + 1; + } + + /// + /// Add a single token to this conversation + /// + /// + /// + /// + /// + public void Prompt(LLamaToken token) + { + AssertCanBePrompted(); + + // Add this token as input + _batchIndex = Executor.Batch.Add(token, _end++, ConversationId, true); + + // Mark this conversation as needing inference/sampling + _requiredEpoch = Executor.Epoch + 1; + } + #endregion +} \ No newline at end of file diff --git a/LLama/Batched/Exceptions.cs b/LLama/Batched/Exceptions.cs new file mode 100644 index 00000000..7d847bfe --- /dev/null +++ b/LLama/Batched/Exceptions.cs @@ -0,0 +1,81 @@ +using System; + +namespace LLama.Batched; + +/// +/// Base class for exceptions thrown from +/// +public class ExperimentalBatchedExecutorException + : Exception +{ + internal ExperimentalBatchedExecutorException(string message) + : base(message) + { + } +} + +/// +/// This exception is thrown when "Prompt()" is called on a which has +/// already been prompted and before "Infer()" has been called on the associated +/// . +/// +public class AlreadyPromptedConversationException + : ExperimentalBatchedExecutorException +{ + internal AlreadyPromptedConversationException() + : base("Must call `Infer()` before prompting this Conversation again") + { + } +} + +/// +/// This exception is thrown when "Sample()" is called on a which has +/// already been prompted and before "Infer()" has been called on the associated +/// . +/// +public class CannotSampleRequiresInferenceException + : ExperimentalBatchedExecutorException +{ + internal CannotSampleRequiresInferenceException() + : base("Must call `Infer()` before sampling from this Conversation") + { + } +} + +/// +/// This exception is thrown when "Sample()" is called on a which was not +/// first prompted. +/// . +/// +public class CannotSampleRequiresPromptException + : ExperimentalBatchedExecutorException +{ + internal CannotSampleRequiresPromptException() + : base("Must call `Prompt()` and then `Infer()` before sampling from this Conversation") + { + } +} + +/// +/// This exception is thrown when "Fork()" is called on a with = true +/// +public class CannotForkWhileRequiresInference + : ExperimentalBatchedExecutorException +{ + internal CannotForkWhileRequiresInference() + : base("Cannot `Fork()` a conversation while RequiresInference is true") + { + } +} + +/// +/// This exception is thrown when "Rewind()" is called on a with = true +/// +public class CannotRewindWhileRequiresInference + : ExperimentalBatchedExecutorException +{ + internal CannotRewindWhileRequiresInference() + : base("Cannot `Rewind()` a conversation while RequiresInference is true") + { + } +} \ No newline at end of file diff --git a/LLama/LLamaContext.cs b/LLama/LLamaContext.cs index 5d026b67..9075c89f 100644 --- a/LLama/LLamaContext.cs +++ b/LLama/LLamaContext.cs @@ -221,7 +221,9 @@ namespace LLama /// The selected token public LLamaToken Sample(ISamplingPipeline pipeline, ReadOnlySpan lastTokens) { - return pipeline.Sample(NativeHandle, NativeHandle.GetLogits(), lastTokens); + var token = pipeline.Sample(NativeHandle, NativeHandle.GetLogits(), lastTokens); + pipeline.Accept(NativeHandle, token); + return token; } /// diff --git a/LLama/LLamaInstructExecutor.cs b/LLama/LLamaInstructExecutor.cs index 969da783..31e975ca 100644 --- a/LLama/LLamaInstructExecutor.cs +++ b/LLama/LLamaInstructExecutor.cs @@ -213,6 +213,7 @@ namespace LLama if (inferenceParams.SamplingPipeline is not null) { id = inferenceParams.SamplingPipeline.Sample(Context.NativeHandle, Context.NativeHandle.GetLogits(), _last_n_tokens.ToArray()); + inferenceParams.SamplingPipeline.Accept(Context.NativeHandle, id); } else { diff --git a/LLama/LLamaInteractExecutor.cs b/LLama/LLamaInteractExecutor.cs index 7d742c81..9338d839 100644 --- a/LLama/LLamaInteractExecutor.cs +++ b/LLama/LLamaInteractExecutor.cs @@ -192,6 +192,7 @@ namespace LLama if (inferenceParams.SamplingPipeline is not null) { id = inferenceParams.SamplingPipeline.Sample(Context.NativeHandle, Context.NativeHandle.GetLogits(), _last_n_tokens.ToArray()); + inferenceParams.SamplingPipeline.Accept(Context.NativeHandle, id); } else { diff --git a/LLama/Native/LLamaBatch.cs b/LLama/Native/LLamaBatch.cs index 532e16ff..50b9c8f0 100644 --- a/LLama/Native/LLamaBatch.cs +++ b/LLama/Native/LLamaBatch.cs @@ -18,6 +18,11 @@ public class LLamaBatch private LLamaSeqId[][] _sequenceIds; private IntPtr[] _sequenceIdsPtrs; + /// + /// Keep track of the index of existing token/position combos in the batch + /// + private readonly Dictionary<(LLamaToken, LLamaPos), int> _index = new(); + /// /// The number of tokens in this batch /// @@ -130,23 +135,44 @@ public class LLamaBatch /// The position to add it att /// The set of sequences to add this token to /// - public void Add(LLamaToken token, LLamaPos pos, ReadOnlySpan sequences, bool logits) + /// The index that the token was added at. Use this for GetLogitsIth + public int Add(LLamaToken token, LLamaPos pos, ReadOnlySpan sequences, bool logits) { + // Try to find this (token, position) combo somewhere in the batch to re-use it + if (_index.TryGetValue((token, pos), out var existingIndex)) + { + if (_sequenceIdCount[existingIndex] + sequences.Length > SequenceCapacity) + GrowMaxSequences(_sequenceIdCount[existingIndex] + sequences.Length); + + foreach (var sequence in sequences) + { + _sequenceIds[existingIndex][_sequenceIdCount[existingIndex]] = sequence; + _sequenceIdCount[existingIndex]++; + } + + return existingIndex; + } + + // Couldn't find this it in the batch, add a new item + + // Frow capacity as necessary if (TokenCount == TokenCapacity) GrowTokenCapacity(); if (sequences.Length > SequenceCapacity) GrowMaxSequences(sequences.Length); + // Store the position in the index, so it can be found later + _index.Add((token, pos), TokenCount); + + // Add the items to the arrays _tokens[TokenCount] = token; _positions[TokenCount] = pos; - _sequenceIdCount[TokenCount] = sequences.Length; for (var i = 0; i < sequences.Length; i++) _sequenceIds[TokenCount][i] = sequences[i]; - _logits[TokenCount] = Convert.ToByte(logits); - TokenCount++; + return TokenCount++; } /// @@ -157,11 +183,12 @@ public class LLamaBatch /// The position to add it att /// The set of sequences to add this token to /// - public void Add(LLamaToken token, LLamaPos pos, List sequences, bool logits) + /// The index that the token was added at. Use this for GetLogitsIth + public int Add(LLamaToken token, LLamaPos pos, List sequences, bool logits) { #if NET5_0_OR_GREATER var seqSpan = CollectionsMarshal.AsSpan(sequences); - Add(token, pos, seqSpan, logits); + return Add(token, pos, seqSpan, logits); #else // on netstandard2.0 we can't use CollectionsMarshal to get directly at the internal memory of // the list. Instead rent an array and copy the data into it. This avoids an allocation, but can't @@ -171,7 +198,7 @@ public class LLamaBatch try { sequences.CopyTo(rented, 0); - Add(token, pos, rented.AsSpan(0, sequences.Count), logits); + return Add(token, pos, rented.AsSpan(0, sequences.Count), logits); } finally { @@ -188,14 +215,15 @@ public class LLamaBatch /// The position to add it att /// The sequence to add this token to /// - public void Add(LLamaToken token, LLamaPos pos, LLamaSeqId sequence, bool logits) + /// The index that the token was added at. Use this for GetLogitsIth + public int Add(LLamaToken token, LLamaPos pos, LLamaSeqId sequence, bool logits) { // Create a temporary span to contain 1 item without allocating Span sequences = stackalloc LLamaSeqId[1]; sequences[0] = sequence; // Add it - Add(token, pos, sequences, logits); + return Add(token, pos, sequences, logits); } /// @@ -205,13 +233,17 @@ public class LLamaBatch /// The starting position to add tokens at /// The sequence to add this token to /// Whether the final token should generate logits - public void AddRange(ReadOnlySpan tokens, LLamaPos start, LLamaSeqId sequence, bool logitsLast) + /// The index that the final token was added at. Use this for GetLogitsIth + public int AddRange(ReadOnlySpan tokens, LLamaPos start, LLamaSeqId sequence, bool logitsLast) { + var last = -1; for (var i = 0; i < tokens.Length; i++) { var logits = (i == tokens.Length - 1) & logitsLast; - Add(tokens[i], start.Value + i, sequence, logits); + last = Add(tokens[i], start.Value + i, sequence, logits); } + + return last; } #endregion @@ -221,5 +253,6 @@ public class LLamaBatch public void Clear() { TokenCount = 0; + _index.Clear(); } } \ No newline at end of file diff --git a/LLama/Sampling/BaseSamplingPipeline.cs b/LLama/Sampling/BaseSamplingPipeline.cs index a41aa67e..b86001aa 100644 --- a/LLama/Sampling/BaseSamplingPipeline.cs +++ b/LLama/Sampling/BaseSamplingPipeline.cs @@ -40,10 +40,7 @@ public abstract class BaseSamplingPipeline var candidates = LLamaTokenDataArray.Create(logits); // Process token data array - ProcessTokenDataArray(ctx, candidates, lastTokens); - - // Choose the final value - return ChooseToken(ctx, candidates); + return ProcessTokenDataArray(ctx, candidates, lastTokens); } finally { @@ -53,6 +50,9 @@ public abstract class BaseSamplingPipeline } } + /// + public abstract void Accept(SafeLLamaContextHandle ctx, LLamaToken token); + #region protected tokens /// /// Get all of the "protected" tokens that cannot be changed by ProcessLogits @@ -107,19 +107,14 @@ public abstract class BaseSamplingPipeline /// protected abstract LLamaToken ProcessTokenDataArray(SafeLLamaContextHandle ctx, LLamaTokenDataArray candidates, ReadOnlySpan lastTokens); - /// - /// Choose the final token from the candidates - /// - /// - /// - /// - protected abstract LLamaToken ChooseToken(SafeLLamaContextHandle ctx, LLamaTokenDataArray candidates); - /// public virtual void Reset() { } + /// + public abstract ISamplingPipeline Clone(); + /// public virtual void Dispose() { diff --git a/LLama/Sampling/DefaultSamplingPipeline.cs b/LLama/Sampling/DefaultSamplingPipeline.cs index b0fb5c59..531f34fa 100644 --- a/LLama/Sampling/DefaultSamplingPipeline.cs +++ b/LLama/Sampling/DefaultSamplingPipeline.cs @@ -141,9 +141,31 @@ public sealed class DefaultSamplingPipeline return id; } + public override void Accept(SafeLLamaContextHandle ctx, LLamaToken token) + { + Grammar?.AcceptToken(ctx, token); + } + /// - protected override LLamaToken ChooseToken(SafeLLamaContextHandle ctx, LLamaTokenDataArray candidates) + public override ISamplingPipeline Clone() { - return candidates.SampleToken(ctx); + var clone = new DefaultSamplingPipeline(); + + foreach (var (k, v) in LogitBias) + clone.LogitBias.Add(k, v); + + clone.Grammar = Grammar?.Clone(); + clone.RepeatPenalty = RepeatPenalty; + clone.AlphaFrequency = AlphaFrequency; + clone.AlphaPresence = AlphaPresence; + clone.Temperature = Temperature; + clone.TopK = TopK; + clone.TailFreeZ = TailFreeZ; + clone.TypicalP = TypicalP; + clone.TopP = TopP; + clone.MinP = MinP; + clone.PenalizeNewline = PenalizeNewline; + + return clone; } } \ No newline at end of file diff --git a/LLama/Sampling/ISamplingPipeline.cs b/LLama/Sampling/ISamplingPipeline.cs index be139879..b538d1fe 100644 --- a/LLama/Sampling/ISamplingPipeline.cs +++ b/LLama/Sampling/ISamplingPipeline.cs @@ -21,10 +21,23 @@ public interface ISamplingPipeline /// LLamaToken Sample(SafeLLamaContextHandle ctx, Span logits, ReadOnlySpan lastTokens); + /// + /// Update the pipeline, with knowledge that a particular token was just accepted + /// + /// + /// + void Accept(SafeLLamaContextHandle ctx, LLamaToken token); + /// /// Reset all internal state of the sampling pipeline /// void Reset(); + + /// + /// Create a copy of this sampling pipeline + /// + /// + ISamplingPipeline Clone(); } ///