Merge pull request #1262 from SignalRT/Optimization

Optimization - A queue with fixed storage size backed by a circular buffer

Author: martindevans

LLama.Benchmark/Collections/FixedSizeQueueBenchmark.cs

@@ -0,0 +1,46 @@
+using System.Linq;
+using BenchmarkDotNet.Attributes;
+using BenchmarkDotNet.Engines;
+using BenchmarkDotNet.Jobs;
+using LLama.Common;
+
+namespace LLama.Benchmark.Collections;
+
+[SimpleJob(RunStrategy.Throughput, RuntimeMoniker.Net80)]
+[MemoryDiagnoser]
+[BenchmarkCategory("Collections", "FixedSizeQueue")]
+public class FixedSizeQueueBenchmark
+{
+ [Params(32, 512, 4096)]
+ public int Capacity { get; set; }
+
+ private int[] _values = Array.Empty<int>();
+
+ [GlobalSetup]
+ public void Setup()
+ {
+ _values = Enumerable.Range(0, Capacity * 4).ToArray();
+ }
+
+ [Benchmark]
+ public int EnqueueWrap()
+ {
+ var queue = new FixedSizeQueue<int>(Capacity);
+ foreach (var value in _values)
+ queue.Enqueue(value);
+ return queue.Count;
+ }
+
+ [Benchmark]
+ public int IterateTailSum()
+ {
+ var queue = new FixedSizeQueue<int>(Capacity);
+ foreach (var value in _values)
+ queue.Enqueue(value);
+
+ var sum = 0;
+ foreach (var value in queue)
+ sum += value;
+ return sum;
+ }
+}

LLama/AntipromptProcessor.cs

@@ -11,7 +11,7 @@ public sealed class AntipromptProcessor
  private int _longestAntiprompt;
  private readonly List<string> _antiprompts = new();
 
- private string? _string;
+ private string _buffer = string.Empty;
 
 
  /// <summary>
@@ -46,6 +46,8 @@ public void SetAntiprompts(IEnumerable<string> antiprompts)
  _longestAntiprompt = 0;
  foreach (var antiprompt in _antiprompts)
  _longestAntiprompt = Math.Max(_longestAntiprompt, antiprompt.Length);
+
+ _buffer = string.Empty;
  }
 
  /// <summary>
@@ -55,21 +57,21 @@ public void SetAntiprompts(IEnumerable<string> antiprompts)
  /// <returns>true if the text buffer ends with any antiprompt</returns>
  public bool Add(string text)
  {
- _string += text;
+ _buffer += text;
 
  // When the string gets very long (4x antiprompt length) trim it down (to 2x antiprompt length).
  // This trimming leaves a lot of extra characters because two sequences can be considered "equal" in unicode
  // even with different numbers of characters. Hopefully there are enough characters here to handle all those weird circumstances!
  var maxLength = Math.Max(32, _longestAntiprompt * 4);
  var trimLength = Math.Max(16, _longestAntiprompt * 2);
- if (_string.Length > maxLength)
- _string = _string.Substring(_string.Length - trimLength);
+ if (_buffer.Length > maxLength)
+ _buffer = _buffer.Substring(_buffer.Length - trimLength);
 
  foreach (var antiprompt in _antiprompts)
- if (_string.EndsWith(antiprompt, StringComparison.CurrentCulture))
+ if (_buffer.EndsWith(antiprompt, StringComparison.CurrentCulture))
  return true;
 
  return false;
  }
  }
-}
+}

LLama/Batched/BatchedExecutor.cs

@@ -1,7 +1,6 @@
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
-using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using LLama.Abstractions;
@@ -16,7 +15,12 @@ public sealed class BatchedExecutor
  : IDisposable
 {
  private int _nextSequenceId;
- private readonly List<IBatch> _batchQueue = [ ];
+ private readonly List<IBatch> _batchQueue = [];
+ private int _batchQueueHead;
+ private int _batchedTokenCount;
+ private bool _batchedTokenCountDirty = true;
+ // Skip compacting the queue until this many processed batches accumulate at the front.
+ private const int CleanupThreshold = 16;
 
  /// <summary>
  /// Set to 1 using interlocked exchange while inference is running
@@ -42,12 +46,27 @@ public sealed class BatchedExecutor
  /// <summary>
  /// Get the number of tokens in the batch, waiting for <see cref="Infer"/> to be called
  /// </summary>
- public int BatchedTokenCount => _batchQueue.Sum(a => a.ItemCount);
+ public int BatchedTokenCount
+ {
+ get
+ {
+ if (_batchedTokenCountDirty)
+ {
+ var total = 0;
+ for (var i = _batchQueueHead; i < _batchQueue.Count; i++)
+ total += _batchQueue[i].ItemCount;
+ _batchedTokenCount = total;
+ _batchedTokenCountDirty = false;
+ }
+
+ return _batchedTokenCount;
+ }
+ }
 
  /// <summary>
  /// Number of batches in the queue, waiting for <see cref="Infer"/> to be called
  /// </summary>
- public int BatchQueueCount => _batchQueue.Count;
+ public int BatchQueueCount => _batchQueue.Count - _batchQueueHead;
 
  /// <summary>
  /// Check if this executor has been disposed.
@@ -147,12 +166,13 @@ public async Task<DecodeResult> Infer(CancellationToken cancellation = default)
  // again after the issue has been fixed (e.g. some KV cache space has been freed) to retry this operation.
  if (status != DecodeResult.Ok)
  {
- _batchQueue.Insert(0, next);
+ RequeueFront(next);
  return status;
  }
 
  // Everything was ok, advance the epoch
  Epoch++;
+ CleanupQueue();
 
  return status;
  }
@@ -166,13 +186,45 @@ public async Task<DecodeResult> Infer(CancellationToken cancellation = default)
 
  IBatch? GetNextBatch()
  {
- if (_batchQueue.Count == 0)
+ if (_batchQueueHead >= _batchQueue.Count)
+ {
+ _batchQueue.Clear();
+ _batchQueueHead = 0;
  return null;
- 
- var nextBatch = _batchQueue[0];
- _batchQueue.RemoveAt(0);
+ }
+
+ var nextBatch = _batchQueue[_batchQueueHead];
+ _batchQueueHead++;
+ _batchedTokenCountDirty = true;
  return nextBatch;
  }
+
+ void RequeueFront(IBatch batch)
+ {
+ Debug.Assert(_batchQueueHead > 0, "Cannot requeue batch when queue head is at zero.");
+ _batchQueue[--_batchQueueHead] = batch;
+ _batchedTokenCountDirty = true;
+ }
+
+ // Remove batches that have already been consumed so the head index does not grow without bound.
+ void CleanupQueue()
+ {
+ if (_batchQueueHead == 0)
+ return;
+
+ if (_batchQueueHead >= _batchQueue.Count)
+ {
+ _batchQueue.Clear();
+ _batchQueueHead = 0;
+ return;
+ }
+
+ if (_batchQueueHead > CleanupThreshold && _batchQueueHead > _batchQueue.Count / 2)
+ {
+ _batchQueue.RemoveRange(0, _batchQueueHead);
+ _batchQueueHead = 0;
+ }
+ }
  }
 
  /// <inheritdoc />
@@ -202,7 +254,7 @@ internal LLamaSeqId GetNextSequenceId()
  throw new ArgumentOutOfRangeException(nameof(minCapacity), $"Request batch capacity must be less than or equal to BatchSize ({Context.BatchSize})");
 
  // Find a batch with space for at least minCapacity tokens
- for (var i = 0; i < _batchQueue.Count; i++)
+ for (var i = _batchQueueHead; i < _batchQueue.Count; i++)
  {
  var item = _batchQueue[i];
  if (item is not TokenBatch { Batch: var batch })
@@ -213,13 +265,17 @@ internal LLamaSeqId GetNextSequenceId()
  continue;
 
  if (batch.TokenCount < Context.BatchSize)
- return (batch, Epoch + (uint)(i + 1) * 2);
+ {
+ _batchedTokenCountDirty = true;
+ return (batch, Epoch + (uint)(i - _batchQueueHead + 1) * 2);
+ }
  }
 
  // Add a new batch to the end of the queue
  var end = new LLamaBatch();
  _batchQueue.Add(new TokenBatch(end));
- return (end, Epoch + (uint)_batchQueue.Count * 2);
+ _batchedTokenCountDirty = true;
+ return (end, Epoch + (uint)(_batchQueue.Count - _batchQueueHead) * 2);
  }
 
  /// <summary>
@@ -234,7 +290,7 @@ internal LLamaSeqId GetNextSequenceId()
  throw new ArgumentOutOfRangeException(nameof(minCapacity), $"Request batch capacity must be less than or equal to BatchSize ({Context.BatchSize})");
 
  // Find a batch with space for at least minCapacity embeddings
- for (var i = 0; i < _batchQueue.Count; i++)
+ for (var i = _batchQueueHead; i < _batchQueue.Count; i++)
  {
  var item = _batchQueue[i];
  if (item is not EmbeddingBatch { Batch: var batch })
@@ -245,13 +301,17 @@ internal LLamaSeqId GetNextSequenceId()
  continue;
 
  if (batch.EmbeddingsCount < Context.BatchSize)
- return (batch, Epoch + (uint)(i + 1) * 2);
+ {
+ _batchedTokenCountDirty = true;
+ return (batch, Epoch + (uint)(i - _batchQueueHead + 1) * 2);
+ }
  }
 
  // Add a new batch to the end of the queue
  var end = new LLamaBatchEmbeddings(Context.EmbeddingSize);
  _batchQueue.Add(new EmbeddingBatch(end));
- return (end, Epoch + (uint)_batchQueue.Count * 2);
+ _batchedTokenCountDirty = true;
+ return (end, Epoch + (uint)(_batchQueue.Count - _batchQueueHead) * 2);
  }
 
  #region batches
@@ -286,4 +346,4 @@ public Task<DecodeResult> DecodeAsync(LLamaContext ctx, CancellationToken token)
  }
  }
  #endregion
-}
+}