Async implementation of LLamaExecutors

LLama/LLamaContext.cs

@@ -564,6 +564,35 @@ public Task<DecodeResult> DecodeAsync(LLamaBatchEmbeddings batch, CancellationTo
         {
             return Task.Run(() => Decode(batch), cancellationToken);
         }
+
+        /// <summary>
+        /// </summary>
+        /// <param name="tokens"></param>
+        /// <param name="id"></param>
+        /// <param name="batch"></param>
+        /// <param name="n_past"></param>
+        /// <returns></returns>
+        public (DecodeResult, int) Decode(List<LLamaToken> tokens, LLamaSeqId id, LLamaBatch batch, ref int n_past)
+        {
+            return NativeHandle.Decode(tokens, id, batch, ref n_past);
+        }
+
+        /// <summary>
+        /// </summary>
+        /// <param name="tokens"></param>
+        /// <param name="id"></param>
+        /// <param name="batch"></param>
+        /// <param name="n_past"></param>
+        /// <returns></returns>
+        public Task<(DecodeResult, int, int)> DecodeAsync(List<LLamaToken> tokens, LLamaSeqId id, LLamaBatch batch, int n_past)
+        {
+            return Task.Run(() =>
+            {
+                var past = n_past;
+                var res = NativeHandle.Decode(tokens, id, batch, ref past);
+                return (res.Item1, res.Item2, past);
+                });
+        }
         #endregion
 
         /// <inheritdoc />

LLama/LLamaInstructExecutor.cs

@@ -106,7 +106,7 @@
             using (var fs = new FileStream(filename, FileMode.Open, FileAccess.Read))
             {
                 var state = await JsonSerializer.DeserializeAsync<InstructExecutorState>(fs);
                 await LoadState(state);
             }
         }
 
@@ -147,11 +147,11 @@
         }
 
         /// <inheritdoc />
         protected override async Task<(bool, IReadOnlyList<string>)> PostProcess(IInferenceParams inferenceParams, InferStateArgs args)
         {
             if (_embed_inps.Count <= _consumedTokensCount)
             {
                 if (_last_n_tokens.TokensEndsWithAnyString(args.Antiprompts, Context.NativeHandle.ModelHandle, Context.Encoding))
                 {
                     args.WaitForInput = true;
                     return (true, Array.Empty<string>());
@@ -177,7 +177,7 @@
         }
 
         /// <inheritdoc />
-        protected override Task InferInternal(IInferenceParams inferenceParams, InferStateArgs args)
+        protected override async Task InferInternal(IInferenceParams inferenceParams, InferStateArgs args)
         {
             var batch = new LLamaBatch();
 
@@ -194,7 +194,9 @@
 
                 TryReuseMatchingPrefix();
 
-                var (result, _) = Context.NativeHandle.Decode(_embeds, LLamaSeqId.Zero, batch, ref _pastTokensCount);
+                var (result, _, pastTokensCount) = await Context.DecodeAsync(_embeds, LLamaSeqId.Zero, batch, _pastTokensCount);
+                _pastTokensCount = pastTokensCount;
+
                 if (result != DecodeResult.Ok)
                     throw new LLamaDecodeError(result);
 
@@ -215,7 +217,7 @@
                 if (!string.IsNullOrEmpty(_pathSession) && args.NeedToSaveSession)
                 {
                     args.NeedToSaveSession = false;
                     SaveSessionFile(_pathSession);
                 }
 
                 LLamaToken id;
@@ -259,7 +261,7 @@
                 }
             }
 
-            return Task.CompletedTask;
+            return;
         }
         /// <summary>
         /// The descriptor of the state of the instruct executor.
@@ -275,7 +277,7 @@
             /// Instruction prefix tokens.
             /// </summary>
             [JsonPropertyName("inp_pfx")]
             public LLamaToken[] InputPrefixTokens { get; set; }
             /// <summary>
             /// Instruction suffix tokens.
             /// </summary>

LLama/LLamaInteractExecutor.cs

@@ -98,7 +98,7 @@
             using (var fs = new FileStream(filename, FileMode.Open, FileAccess.Read))
             {
                 var state = await JsonSerializer.DeserializeAsync<InteractiveExecutorState>(fs);
                 await LoadState(state);
             }
         }
 
@@ -159,7 +159,7 @@
             {
                 foreach (var image in Images)
                 {
                     _imageEmbedHandles.Add(SafeLlavaImageEmbedHandle.CreateFromMemory(ClipModel.NativeHandle, Context, image));
                 }
 
                 int imageIndex = text.IndexOf("<image>");
@@ -196,11 +196,11 @@
         /// <param name="inferenceParams"></param>
         /// <param name="args"></param>
         /// <returns></returns>
         protected override async Task<(bool, IReadOnlyList<string>)> PostProcess(IInferenceParams inferenceParams, InferStateArgs args)
         {
             if (_embed_inps.Count <= _consumedTokensCount)
             {
                 if (_last_n_tokens.TokensEndsWithAnyString(args.Antiprompts, Context.NativeHandle.ModelHandle, Context.Encoding))
                     args.WaitForInput = true;
 
                 if (_pastTokensCount > 0 && args.WaitForInput)
@@ -222,7 +222,7 @@
         }
 
         /// <inheritdoc />
-        protected override Task InferInternal(IInferenceParams inferenceParams, InferStateArgs args)
+        protected override async Task InferInternal(IInferenceParams inferenceParams, InferStateArgs args)
         {
             var batch = new LLamaBatch();
 
@@ -250,27 +250,32 @@
 
                 // Changes to support Multi-Modal LLMs.
                 //
-                (DecodeResult, int) header, end, result;
+                (DecodeResult, int, int) header, end, result;
                 if (IsMultiModal &&  _EmbedImagePosition > 0)
                 {
                     // Tokens previous to the images
-                    header = Context.NativeHandle.Decode(_embeds.GetRange(0, _EmbedImagePosition), LLamaSeqId.Zero, batch, ref _pastTokensCount);
+                    header = await Context.DecodeAsync(_embeds.GetRange(0, _EmbedImagePosition), LLamaSeqId.Zero, batch, _pastTokensCount);
+                    _pastTokensCount = header.Item3;
+
                     if (header.Item1 != DecodeResult.Ok) throw new LLamaDecodeError(header.Item1);
 
                     // Images
                     foreach( var image in _imageEmbedHandles )
                         ClipModel.EvalImageEmbed(Context, image, ref _pastTokensCount);
 
                     // Post-image Tokens
-                    end = Context.NativeHandle.Decode(_embeds.GetRange(_EmbedImagePosition, _embeds.Count - _EmbedImagePosition), LLamaSeqId.Zero, batch, ref _pastTokensCount);
+                    end = await Context.DecodeAsync(_embeds.GetRange(_EmbedImagePosition, _embeds.Count - _EmbedImagePosition), LLamaSeqId.Zero, batch, _pastTokensCount);
+                    _pastTokensCount = end.Item3;
 
                     _EmbedImagePosition = -1;
                     _imageEmbedHandles.Clear();
                     Images.Clear();
                 }
                 else
                 {
-                    result = Context.NativeHandle.Decode(_embeds, LLamaSeqId.Zero, batch, ref _pastTokensCount);
+                    result = await Context.DecodeAsync(_embeds, LLamaSeqId.Zero, batch, _pastTokensCount);
+                    _pastTokensCount = result.Item3;
+
                     if (result.Item1 != DecodeResult.Ok) throw new LLamaDecodeError(result.Item1);
                 }
 
@@ -346,7 +351,7 @@
                 }
             }
 
-            return Task.CompletedTask;
+            return;
         }
 
         /// <summary>

LLama/LLamaStatelessExecutor.cs

@@ -96,7 +96,9 @@ public async IAsyncEnumerable<string> InferAsync(string prompt, IInferenceParams
 
             // Evaluate the prompt, in chunks smaller than the max batch size
             var n_past = 0;
-            var (r, _) = Context.NativeHandle.Decode(tokens, LLamaSeqId.Zero, _batch, ref n_past);
+            var (r, _, past) = await Context.DecodeAsync(tokens, LLamaSeqId.Zero, _batch, n_past);
+            n_past = past;
+
             if (r != DecodeResult.Ok)
                 throw new LLamaDecodeError(r);