From f6d15148e9e7b494b04931f1c948ef5a218b5d91 Mon Sep 17 00:00:00 2001 From: james-333i Date: Tue, 25 Aug 2026 14:36:42 -0700 Subject: [PATCH 1/3] Add vision support to LlamaLanguageModel via mtmd Image segments threw unsupportedFeature because the backend had no multimodal path, even though the prebuilt llama.cpp binaries ship the mtmd library and its helpers. Accept an mmprojPath at initialization and load the projector next to the model. When a projector is present, prompt formatting replaces each image segment with the mtmd media marker and collects payloads in order, then generation tokenizes the marker-annotated prompt with mtmd_tokenize and evaluates text and image chunks through mtmd_helper_eval_chunks before sampling continues from the resulting position. Both respond and streaming support images, and models without a projector keep rejecting image input. Adds live tests generating from an embedded test image through both paths. --- .../Models/LlamaLanguageModel.swift | 353 ++++++++++++++++-- .../LlamaLanguageModelTests.swift | 75 ++++ 2 files changed, 397 insertions(+), 31 deletions(-) diff --git a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift index d2f29b03..4dd41aa0 100644 --- a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift @@ -208,6 +208,11 @@ import Foundation /// A negative value offloads all layers, and `0` runs entirely on the CPU. public let gpuLayers: Int32 + /// The path to the multimodal projector GGUF file, when the model has one. + /// + /// Prompts may include image segments only when a projector is loaded. + public let mmprojPath: String? + /// The default GPU layer count for the current platform. /// /// All layers are offloaded by default: the prebuilt llama.cpp binaries @@ -453,6 +458,9 @@ import Foundation /// The model's vocabulary private var vocab: OpaquePointer? + /// The multimodal projector context, when a projector file was provided + private var mtmdContext: OpaquePointer? + /// Whether the model is currently loaded private var isModelLoaded: Bool = false @@ -462,9 +470,18 @@ import Foundation /// - modelPath: The path to the GGUF model file. /// - gpuLayers: The number of model layers to offload to the GPU. /// Defaults to ``defaultGPULayerCount``. - public init(modelPath: String, gpuLayers: Int32 = LlamaLanguageModel.defaultGPULayerCount) { + /// - mmprojPath: The path to a multimodal projector GGUF file matching + /// the model. When provided, prompts may include image segments, + /// which are encoded through the projector. Defaults to `nil` + /// (text only). + public init( + modelPath: String, + gpuLayers: Int32 = LlamaLanguageModel.defaultGPULayerCount, + mmprojPath: String? = nil + ) { self.modelPath = modelPath self.gpuLayers = gpuLayers + self.mmprojPath = mmprojPath self.legacyDefaults = ResolvedGenerationOptions() } @@ -505,6 +522,9 @@ import Foundation } deinit { + if let mtmdContext = mtmdContext { + mtmd_free(mtmdContext) + } if let model = model { llama_model_free(model) } @@ -517,8 +537,9 @@ import Foundation includeSchemaInPrompt: Bool, options: GenerationOptions ) async throws -> LanguageModelSession.Response where Content: Generable { - // Validate that no image segments are present - try validateNoImageSegments(in: session) + if mmprojPath == nil { + try validateNoImageSegments(in: session) + } try await ensureModelLoaded() let runtimeOptions = resolvedOptions(from: options) @@ -542,6 +563,9 @@ import Foundation llama_set_causal_attn(context, true) llama_set_n_threads(context, runtimeOptions.threads, runtimeOptions.threads) + var promptImages: [Data] = [] + let imageMarker = mtmdContext != nil ? String(cString: mtmd_default_marker()) : nil + let fullPrompt: String if includeSchemaInPrompt, type != String.self { fullPrompt = try formatPrompt( @@ -550,18 +574,39 @@ import Foundation assistantPrefill: runtimeOptions.assistantPrefill ) } else { - fullPrompt = try formatPrompt(for: session, assistantPrefill: runtimeOptions.assistantPrefill) + fullPrompt = try formatPrompt( + for: session, + extraSystemMessage: nil, + assistantPrefill: runtimeOptions.assistantPrefill, + imageMarker: imageMarker, + images: &promptImages + ) } if type == String.self { let maxTokens = runtimeOptions.maximumResponseTokens ?? 100 - let text = try await generateText( - context: context, - model: model!, - prompt: fullPrompt, - maxTokens: maxTokens, - options: runtimeOptions - ) + let text: String + if promptImages.isEmpty { + text = try await generateText( + context: context, + model: model!, + prompt: fullPrompt, + maxTokens: maxTokens, + options: runtimeOptions + ) + } else { + var accumulated = "" + try performMultimodalGeneration( + context: context, + prompt: fullPrompt, + images: promptImages, + maxTokens: maxTokens, + options: runtimeOptions + ) { tokenText in + accumulated += tokenText + } + text = accumulated + } return LanguageModelSession.Response( content: text as! Content, @@ -599,15 +644,16 @@ import Foundation fatalError("LlamaLanguageModel only supports generating String content") } - // Validate that no image segments are present - do { - try validateNoImageSegments(in: session) - } catch { - return LanguageModelSession.ResponseStream( - stream: AsyncThrowingStream { continuation in - continuation.finish(throwing: error) - } - ) + if mmprojPath == nil { + do { + try validateNoImageSegments(in: session) + } catch { + return LanguageModelSession.ResponseStream( + stream: AsyncThrowingStream { continuation in + continuation.finish(throwing: error) + } + ) + } } let stream: AsyncThrowingStream.Snapshot, any Error> = @@ -637,18 +683,18 @@ import Foundation llama_set_n_threads(context, runtimeOptions.threads, runtimeOptions.threads) var accumulatedText = "" + var promptImages: [Data] = [] + let imageMarker = + self.mtmdContext != nil ? String(cString: mtmd_default_marker()) : nil let fullPrompt = try self.formatPrompt( for: session, - assistantPrefill: runtimeOptions.assistantPrefill + extraSystemMessage: nil, + assistantPrefill: runtimeOptions.assistantPrefill, + imageMarker: imageMarker, + images: &promptImages ) - try self.performTextGeneration( - context: context, - model: model!, - prompt: fullPrompt, - maxTokens: maxTokens, - options: runtimeOptions - ) { tokenText in + let yieldToken: (String) -> Void = { tokenText in accumulatedText += tokenText let snapshot = LanguageModelSession.ResponseStream.Snapshot( @@ -658,6 +704,26 @@ import Foundation continuation.yield(snapshot) } + if promptImages.isEmpty { + try self.performTextGeneration( + context: context, + model: model!, + prompt: fullPrompt, + maxTokens: maxTokens, + options: runtimeOptions, + onToken: yieldToken + ) + } else { + try self.performMultimodalGeneration( + context: context, + prompt: fullPrompt, + images: promptImages, + maxTokens: maxTokens, + options: runtimeOptions, + onToken: yieldToken + ) + } + continuation.finish() } catch { continuation.finish(throwing: error) @@ -686,6 +752,10 @@ import Foundation llama_backend_init() // Free any existing model before loading a new one + if let existingContext = mtmdContext { + mtmd_free(existingContext) + self.mtmdContext = nil + } if let existingModel = model { llama_model_free(existingModel) self.model = nil @@ -696,6 +766,22 @@ import Foundation throw LlamaLanguageModelError.modelLoadFailed } + if let mmprojPath { + guard FileManager.default.fileExists(atPath: mmprojPath) else { + llama_model_free(loadedModel) + throw LlamaLanguageModelError.invalidModelPath + } + var mtmdParams = mtmd_context_params_default() + mtmdParams.use_gpu = gpuLayers != 0 + mtmdParams.print_timings = false + mtmdParams.n_threads = legacyDefaults.threads + guard let projector = mtmd_init_from_file(mmprojPath, loadedModel, mtmdParams) else { + llama_model_free(loadedModel) + throw LlamaLanguageModelError.modelLoadFailed + } + self.mtmdContext = projector + } + self.model = loadedModel self.vocab = llama_model_get_vocab(loadedModel) self.isModelLoaded = true @@ -1285,6 +1371,146 @@ import Foundation } } + /// Evaluates a marker-annotated multimodal prompt through the projector, + /// then generates text tokens from the resulting state. + private func performMultimodalGeneration( + context: OpaquePointer, + prompt: String, + images: [Data], + maxTokens: Int, + options: ResolvedGenerationOptions, + onToken: (String) -> Void + ) throws { + guard let mtmdContext, let model = self.model, + let vocab = llama_model_get_vocab(model) + else { + throw LlamaLanguageModelError.contextInitializationFailed + } + + var bitmaps: [OpaquePointer?] = [] + defer { + for bitmap in bitmaps { + if let bitmap { + mtmd_bitmap_free(bitmap) + } + } + } + for imageData in images { + // Pinned to the current llama.swift signature. llama.cpp master adds a + // trailing options argument to this helper; update alongside the dependency. + let wrapper = imageData.withUnsafeBytes { raw -> mtmd_helper_bitmap_wrapper in + mtmd_helper_bitmap_init_from_buf( + mtmdContext, + raw.bindMemory(to: UInt8.self).baseAddress, + imageData.count, + false + ) + } + if let videoContext = wrapper.video_ctx { + mtmd_helper_video_free(videoContext) + throw LlamaLanguageModelError.unsupportedFeature + } + guard let bitmap = wrapper.bitmap else { + throw LlamaLanguageModelError.encodingFailed + } + bitmaps.append(bitmap) + } + + guard let chunks = mtmd_input_chunks_init() else { + throw LlamaLanguageModelError.encodingFailed + } + defer { mtmd_input_chunks_free(chunks) } + + let tokenizeResult = prompt.withCString { cPrompt -> Int32 in + var inputText = mtmd_input_text( + text: cPrompt, + text_len: strlen(cPrompt), + add_special: true, + parse_special: true + ) + return bitmaps.withUnsafeMutableBufferPointer { buffer in + mtmd_tokenize(mtmdContext, chunks, &inputText, buffer.baseAddress, buffer.count) + } + } + guard tokenizeResult == 0 else { + throw LlamaLanguageModelError.tokenizationFailed + } + + var pastPosition: llama_pos = 0 + let evalResult = mtmd_helper_eval_chunks( + mtmdContext, + context, + chunks, + 0, + 0, + Int32(options.batchSize), + true, + &pastPosition + ) + guard evalResult == 0 else { + throw LlamaLanguageModelError.decodingFailed + } + + guard let sampler = llama_sampler_chain_init(llama_sampler_chain_default_params()) else { + throw LlamaLanguageModelError.decodingFailed + } + defer { llama_sampler_free(sampler) } + let samplerPtr = UnsafeMutablePointer(sampler) + + if options.repeatPenalty != 1.0 || options.frequencyPenalty != 0.0 || options.presencePenalty != 0.0 { + llama_sampler_chain_add( + samplerPtr, + llama_sampler_init_penalties( + llama_vocab_n_tokens(vocab), + options.repeatLastN, + options.repeatPenalty, + options.frequencyPenalty, + options.presencePenalty + ) + ) + } + applySampling(sampler: samplerPtr, effectiveTemperature: options.temperature, options: options) + + var batch = llama_batch_init(1, 0, 1) + defer { llama_batch_free(batch) } + + var n_cur: Int32 = Int32(pastPosition) + var sampleIndex: Int32 = -1 + + for _ in 0 ..< maxTokens { + if Task.isCancelled { + break + } + + let nextToken = llama_sampler_sample(samplerPtr, context, sampleIndex) + llama_sampler_accept(samplerPtr, nextToken) + + if llama_vocab_is_eog(vocab, nextToken) { + break + } + + if let tokenText = tokenToText(vocab: vocab, token: nextToken) { + onToken(tokenText) + } + + batch.n_tokens = 1 + batch.token[0] = nextToken + batch.pos[0] = n_cur + batch.n_seq_id[0] = 1 + if let seq_ids = batch.seq_id, let seq_id = seq_ids[0] { + seq_id[0] = 0 + } + batch.logits[0] = 1 + + n_cur += 1 + + guard llama_decode(context, batch) == 0 else { + break + } + sampleIndex = 0 + } + } + // MARK: - Image Validation private func validateNoImageSegments(in session: LanguageModelSession) throws { @@ -1422,6 +1648,23 @@ import Foundation for session: LanguageModelSession, extraSystemMessage: String? = nil, assistantPrefill: String? = nil + ) throws -> String { + var images: [Data] = [] + return try formatPrompt( + for: session, + extraSystemMessage: extraSystemMessage, + assistantPrefill: assistantPrefill, + imageMarker: nil, + images: &images + ) + } + + private func formatPrompt( + for session: LanguageModelSession, + extraSystemMessage: String?, + assistantPrefill: String?, + imageMarker: String?, + images: inout [Data] ) throws -> String { guard let model = self.model else { throw LlamaLanguageModelError.modelLoadFailed @@ -1432,19 +1675,31 @@ import Foundation for entry in session.transcript { switch entry { case .instructions(let instructions): - let text = extractText(from: instructions.segments) + let text = try extractContent( + from: instructions.segments, + imageMarker: imageMarker, + images: &images + ) if !text.isEmpty { messages.append(("system", text)) } case .prompt(let prompt): - let text = extractText(from: prompt.segments) + let text = try extractContent( + from: prompt.segments, + imageMarker: imageMarker, + images: &images + ) if !text.isEmpty { messages.append(("user", text)) } case .response(let response): - let text = extractText(from: response.segments) + let text = try extractContent( + from: response.segments, + imageMarker: imageMarker, + images: &images + ) if !text.isEmpty { messages.append(("assistant", text)) } @@ -1522,6 +1777,42 @@ import Foundation }.joined() } + /// Extracts message content from segments, replacing each image segment + /// with `imageMarker` and collecting its payload in order. Image segments + /// throw ``LlamaLanguageModelError/unsupportedFeature`` when no marker is + /// provided. + private func extractContent( + from segments: [Transcript.Segment], + imageMarker: String?, + images: inout [Data] + ) throws -> String { + var parts: [String] = [] + for segment in segments { + switch segment { + case .text(let t): + parts.append(t.content) + case .image(let image): + guard let imageMarker else { + throw LlamaLanguageModelError.unsupportedFeature + } + switch image.source { + case .data(let data, _): + images.append(data) + parts.append(imageMarker) + case .url(let url): + guard url.isFileURL, let data = try? Data(contentsOf: url) else { + throw LlamaLanguageModelError.unsupportedFeature + } + images.append(data) + parts.append(imageMarker) + } + default: + break + } + } + return parts.joined() + } + private func tokenizeText(vocab: OpaquePointer, text: String) throws -> [llama_token] { let utf8Count = text.utf8.count let maxTokens = Int32(max(utf8Count * 2, 8)) // Rough estimate, minimum capacity diff --git a/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift index a319b82c..ed5a18f6 100644 --- a/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift @@ -425,3 +425,78 @@ import Testing } } #endif // Llama + +#if Llama + @Suite( + "LlamaLanguageModel vision", + .serialized, + .enabled( + if: ProcessInfo.processInfo.environment["LLAMA_VISION_MODEL_PATH"] != nil + && ProcessInfo.processInfo.environment["LLAMA_VISION_MMPROJ_PATH"] != nil + ) + ) + struct LlamaLanguageModelVisionTests { + static let redSquarePNG = Data( + base64Encoded: "iVBORw0KGgoAAAANSUhEUgAAAGAAAABgCAIAAABt+uBvAAABC0lEQVR4nO3OMQ0AIAAEsfdvGhyw9gaS" + + "CujO9j34QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E" + + "+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E" + + "+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E" + + "+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E" + + "+UHcBWwZ3g5gacwjAAAAAElFTkSuQmCC" + )! + + let model = LlamaLanguageModel( + modelPath: ProcessInfo.processInfo.environment["LLAMA_VISION_MODEL_PATH"]!, + mmprojPath: ProcessInfo.processInfo.environment["LLAMA_VISION_MMPROJ_PATH"]! + ) + + @Test func describesImageData() async throws { + let transcript = Transcript(entries: [ + .prompt( + Transcript.Prompt(segments: [ + .text(.init(content: "What is the dominant color of this image? Answer with one word.")), + .image(.init(data: Self.redSquarePNG, mimeType: "image/png")), + ]) + ) + ]) + let session = LanguageModelSession(model: model, transcript: transcript) + let response = try await session.respond(to: "") + #expect(response.content.lowercased().contains("red")) + } + + @Test func streamsImageDescription() async throws { + let transcript = Transcript(entries: [ + .prompt( + Transcript.Prompt(segments: [ + .text(.init(content: "What is the dominant color of this image? Answer with one word.")), + .image(.init(data: Self.redSquarePNG, mimeType: "image/png")), + ]) + ) + ]) + let session = LanguageModelSession(model: model, transcript: transcript) + let stream = session.streamResponse(to: "") + var last = "" + for try await snapshot in stream { + last = snapshot.content + } + #expect(last.lowercased().contains("red")) + } + + @Test func rejectsImagesWithoutProjector() async throws { + let textOnlyModel = LlamaLanguageModel( + modelPath: ProcessInfo.processInfo.environment["LLAMA_VISION_MODEL_PATH"]! + ) + let transcript = Transcript(entries: [ + .prompt( + Transcript.Prompt(segments: [ + .image(.init(data: Self.redSquarePNG, mimeType: "image/png")) + ]) + ) + ]) + let session = LanguageModelSession(model: textOnlyModel, transcript: transcript) + await #expect(throws: LlamaLanguageModelError.self) { + _ = try await session.respond(to: "") + } + } + } +#endif From 03ada51f829675e59f4aa6629f4c13624374119d Mon Sep 17 00:00:00 2001 From: james-333i Date: Tue, 25 Aug 2026 15:13:49 -0700 Subject: [PATCH 2/3] Render the Gemma 4 chat format when template detection fails Gemma 4's canonical chat template no longer contains the start_of_turn marker that llama_chat_apply_template keys its Gemma detection on, so formatting threw encodingFailed for every Gemma 4 GGUF. When template application fails and the embedded template carries the Gemma 4 turn syntax, render it directly: turns open with a turn marker and role, close with the reverse marker, the assistant role is named model, and generation opens a model turn. The BOS token is applied during tokenization, and thinking is opt-in in this format so no suppression is needed. --- .../Models/LlamaLanguageModel.swift | 24 ++++++++++++ .../LlamaGemma4TemplateTests.swift | 37 +++++++++++++++++++ 2 files changed, 61 insertions(+) create mode 100644 Tests/AnyLanguageModelTests/LlamaGemma4TemplateTests.swift diff --git a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift index 4dd41aa0..bc5d3340 100644 --- a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift @@ -1741,6 +1741,9 @@ import Foundation ) guard requiredSize > 0 else { + if let tmpl, String(cString: tmpl).contains("<|turn>") { + return Self.renderGemma4Prompt(messages: messages, assistantPrefill: assistantPrefill) + } throw LlamaLanguageModelError.encodingFailed } @@ -1770,6 +1773,27 @@ import Foundation return rendered } + /// Renders the Gemma 4 canonical chat format, which + /// `llama_chat_apply_template` does not recognize: turns open with + /// `<|turn>role`, close with ``, and the assistant role is named + /// `model`. The BOS token is applied during tokenization. + static func renderGemma4Prompt( + messages: [(role: String, content: String)], + assistantPrefill: String? + ) -> String { + var rendered = "" + for message in messages { + let role = message.role == "assistant" ? "model" : message.role + let content = message.content.trimmingCharacters(in: .whitespacesAndNewlines) + rendered += "<|turn>\(role)\n\(content)\n" + } + rendered += "<|turn>model\n" + if let assistantPrefill, !assistantPrefill.isEmpty { + rendered += assistantPrefill + } + return rendered + } + private func extractText(from segments: [Transcript.Segment]) -> String { segments.compactMap { segment -> String? in if case .text(let t) = segment { return t.content } diff --git a/Tests/AnyLanguageModelTests/LlamaGemma4TemplateTests.swift b/Tests/AnyLanguageModelTests/LlamaGemma4TemplateTests.swift new file mode 100644 index 00000000..756eb412 --- /dev/null +++ b/Tests/AnyLanguageModelTests/LlamaGemma4TemplateTests.swift @@ -0,0 +1,37 @@ +import Foundation +import Testing + +@testable import AnyLanguageModel + +#if Llama + @Suite("LlamaLanguageModel Gemma 4 template") + struct LlamaGemma4TemplateTests { + @Test func rendersTurnsAndOpensTheModelTurn() { + let rendered = LlamaLanguageModel.renderGemma4Prompt( + messages: [ + (role: "system", content: "Be brief."), + (role: "user", content: "Hi"), + (role: "assistant", content: "Hello"), + (role: "user", content: "Bye"), + ], + assistantPrefill: nil + ) + #expect( + rendered + == "<|turn>system\nBe brief.\n" + + "<|turn>user\nHi\n" + + "<|turn>model\nHello\n" + + "<|turn>user\nBye\n" + + "<|turn>model\n" + ) + } + + @Test func appendsTheAssistantPrefill() { + let rendered = LlamaLanguageModel.renderGemma4Prompt( + messages: [(role: "user", content: "Hi")], + assistantPrefill: "" + ) + #expect(rendered == "<|turn>user\nHi\n<|turn>model\n") + } + } +#endif From e79f8f1019f0ff129e928572ac9f082fd5eec4ab Mon Sep 17 00:00:00 2001 From: james-333i Date: Tue, 25 Aug 2026 20:09:26 -0700 Subject: [PATCH 3/3] Reuse a per-session context across chat exchanges Every generation created a fresh llama_context and prefilled the full rendered conversation from token zero, so multi-turn chat cost grew with the square of the transcript and long conversations spent most of their time re-decoding history. Keep one context alive per session for plain chat generations. Each exchange tokenizes the rendered prompt, keeps the longest token prefix shared with the context's recorded state, removes diverged state with llama_memory_seq_rm, and decodes only the remainder. Backends that cannot rewind, such as recurrent models, fall back to clearing memory and decoding the full prompt, and appends need no rewind on any backend. The final prompt token is always re-decoded so sampling has fresh logits, generated tokens extend the recorded state as they decode, and any generation error discards the cached context. Structured generation, image prompts, and encoder models keep single-use contexts, and clearCachedContext lets consumers free the cached state under memory pressure. Adds a live test asserting prefix reuse on the second turn of a session. --- .../Models/LlamaLanguageModel.swift | 429 +++++++++++------- .../LlamaLanguageModelTests.swift | 20 + 2 files changed, 289 insertions(+), 160 deletions(-) diff --git a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift index bc5d3340..e6204a1f 100644 --- a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift @@ -464,6 +464,234 @@ import Foundation /// Whether the model is currently loaded private var isModelLoaded: Bool = false + /// A context kept alive for one session so exchanges reuse its state. + private struct CachedSessionContext { + let sessionID: ObjectIdentifier + let context: OpaquePointer + var tokens: [llama_token] + let contextSize: UInt32 + let batchSize: UInt32 + /// Whether a generation is currently decoding on the context. + var isCheckedOut: Bool + /// Whether the context should be freed once the current generation releases it. + var discardWhenReleased: Bool + } + + /// Guards `cachedSessionContext`. A generation checks the cached context out + /// for its whole run, so a concurrent generation for another session never + /// frees a context that is still decoding: it runs on a transient context + /// instead and leaves the cache untouched. + private let sessionContextLock = NSLock() + private var cachedSessionContext: CachedSessionContext? + + /// The number of prompt tokens reused from the cached context by the most + /// recent chat generation. + internal private(set) var lastReusedTokenCount: Int = 0 + + /// The number of prompt tokens decoded by the most recent chat generation. + internal private(set) var lastPrefillTokenCount: Int = 0 + + /// Frees the cached per-session context and the state it holds. + /// + /// The cached context, including its KV state, otherwise lives as long as the + /// model. The next chat generation prefills its full prompt again. Call this + /// under memory pressure or when a session is discarded. If a generation is + /// running on the cached context, it is freed as soon as that generation ends. + public func clearCachedContext() { + discardCachedSessionContext() + } + + private func discardCachedSessionContext() { + sessionContextLock.lock() + defer { sessionContextLock.unlock() } + guard var cached = cachedSessionContext else { return } + if cached.isCheckedOut { + cached.discardWhenReleased = true + cachedSessionContext = cached + return + } + llama_free(cached.context) + cachedSessionContext = nil + } + + private func recordCachedTokens(_ tokens: [llama_token], context: OpaquePointer) { + sessionContextLock.lock() + defer { sessionContextLock.unlock() } + guard var cached = cachedSessionContext, cached.context == context else { return } + cached.tokens = tokens + cachedSessionContext = cached + } + + /// Returns a context obtained from `acquireSessionContext`. The cached + /// context is checked back in (or freed, if a discard was requested while + /// it was busy); a transient context is freed. + private func releaseSessionContext(_ context: OpaquePointer) { + sessionContextLock.lock() + defer { sessionContextLock.unlock() } + if var cached = cachedSessionContext, cached.context == context { + if cached.discardWhenReleased { + llama_free(cached.context) + cachedSessionContext = nil + } else { + cached.isCheckedOut = false + cachedSessionContext = cached + } + return + } + llama_free(context) + } + + /// Returns a context for the session along with the index of the first + /// prompt token that still needs to be decoded. Pair every call with + /// `releaseSessionContext(_:)` once generation ends. + /// + /// A cached context whose recorded tokens share a prefix with the prompt + /// keeps that prefix: matching state past the divergence point is removed + /// with `llama_memory_seq_rm`, and backends whose state cannot be rewound + /// (recurrent models) fall back to clearing the memory and decoding the + /// full prompt. The final prompt token is always re-decoded so sampling + /// has fresh logits. + /// + /// While another generation holds the cached context, the caller gets a + /// transient context that decodes the full prompt and is not cached. + private func acquireSessionContext( + for session: LanguageModelSession, + promptTokens: [llama_token], + options: ResolvedGenerationOptions + ) throws -> (context: OpaquePointer, startIndex: Int) { + let sessionID = ObjectIdentifier(session) + sessionContextLock.lock() + defer { sessionContextLock.unlock() } + + if var cached = cachedSessionContext, + !cached.isCheckedOut, + cached.sessionID == sessionID, + cached.contextSize == options.contextSize, + cached.batchSize == options.batchSize + { + var common = 0 + while common < cached.tokens.count, common < promptTokens.count, + cached.tokens[common] == promptTokens[common] + { + common += 1 + } + if common == promptTokens.count { + common = max(0, promptTokens.count - 1) + } + if common < cached.tokens.count { + let memory = llama_get_memory(cached.context) + if !llama_memory_seq_rm(memory, 0, llama_pos(common), -1) { + llama_memory_clear(memory, true) + common = 0 + } + } + cached.tokens = Array(promptTokens.prefix(common)) + cached.isCheckedOut = true + cachedSessionContext = cached + return (cached.context, common) + } + + if let cached = cachedSessionContext, cached.isCheckedOut { + return (try makeFreshContext(options: options), 0) + } + + if let cached = cachedSessionContext { + llama_free(cached.context) + cachedSessionContext = nil + } + let contextParams = createContextParams(from: options) + guard let context = llama_init_from_model(model!, contextParams) else { + throw LlamaLanguageModelError.contextInitializationFailed + } + guard llama_get_memory(context) != nil else { + llama_free(context) + throw LlamaLanguageModelError.encoderOnlyModel + } + cachedSessionContext = CachedSessionContext( + sessionID: sessionID, + context: context, + tokens: [], + contextSize: options.contextSize, + batchSize: options.batchSize, + isCheckedOut: true, + discardWhenReleased: false + ) + return (context, 0) + } + + /// Creates a single-use context for generations that do not reuse state. + private func makeFreshContext(options: ResolvedGenerationOptions) throws -> OpaquePointer { + let contextParams = createContextParams(from: options) + guard let context = llama_init_from_model(model!, contextParams) else { + throw LlamaLanguageModelError.contextInitializationFailed + } + guard llama_get_memory(context) != nil else { + llama_free(context) + throw LlamaLanguageModelError.encoderOnlyModel + } + llama_set_causal_attn(context, true) + llama_set_n_threads(context, options.threads, options.threads) + return context + } + + /// Runs a chat text generation for the session, reusing the session's + /// cached context when its state matches a prefix of the prompt. + private func generateChatText( + session: LanguageModelSession, + prompt: String, + maxTokens: Int, + options: ResolvedGenerationOptions, + onToken: (String) -> Void + ) throws { + guard let model = self.model, let vocab = llama_model_get_vocab(model) else { + throw LlamaLanguageModelError.contextInitializationFailed + } + + let promptTokens = try tokenizeText(vocab: vocab, text: prompt) + guard !promptTokens.isEmpty else { + throw LlamaLanguageModelError.tokenizationFailed + } + + if llama_model_has_encoder(model) { + let context = try makeFreshContext(options: options) + defer { llama_free(context) } + try performTokenGeneration( + context: context, + vocab: vocab, + promptTokens: promptTokens, + startIndex: 0, + maxTokens: maxTokens, + options: options, + onToken: onToken + ) + return + } + + let (context, startIndex) = try acquireSessionContext( + for: session, + promptTokens: promptTokens, + options: options + ) + defer { releaseSessionContext(context) } + llama_set_causal_attn(context, true) + llama_set_n_threads(context, options.threads, options.threads) + + do { + try performTokenGeneration( + context: context, + vocab: vocab, + promptTokens: promptTokens, + startIndex: startIndex, + maxTokens: maxTokens, + options: options, + onToken: onToken + ) + } catch { + discardCachedSessionContext() + throw error + } + } + /// Creates a Llama language model. /// /// - Parameters: @@ -522,6 +750,9 @@ import Foundation } deinit { + if let cached = cachedSessionContext { + llama_free(cached.context) + } if let mtmdContext = mtmdContext { mtmd_free(mtmdContext) } @@ -544,24 +775,6 @@ import Foundation let runtimeOptions = resolvedOptions(from: options) let structuredOptions = resolvedStructuredOptions(from: options) - let contextParams = createContextParams(from: runtimeOptions) - - // Try to create context with error handling - guard let context = llama_init_from_model(model!, contextParams) else { - throw LlamaLanguageModelError.contextInitializationFailed - } - defer { llama_free(context) } - - // Check if this is an embedding model (no KV cache). - // This early check catches models configured for embeddings that lack a KV cache. - // A complementary architectural check in prepareInitialBatch catches encoder-only - // models (like BERT) by their architecture type. - if llama_get_memory(context) == nil { - throw LlamaLanguageModelError.encoderOnlyModel - } - - llama_set_causal_attn(context, true) - llama_set_n_threads(context, runtimeOptions.threads, runtimeOptions.threads) var promptImages: [Data] = [] let imageMarker = mtmdContext != nil ? String(cString: mtmd_default_marker()) : nil @@ -587,14 +800,20 @@ import Foundation let maxTokens = runtimeOptions.maximumResponseTokens ?? 100 let text: String if promptImages.isEmpty { - text = try await generateText( - context: context, - model: model!, + var accumulated = "" + try generateChatText( + session: session, prompt: fullPrompt, maxTokens: maxTokens, options: runtimeOptions - ) + ) { tokenText in + accumulated += tokenText + } + text = accumulated } else { + discardCachedSessionContext() + let context = try makeFreshContext(options: runtimeOptions) + defer { llama_free(context) } var accumulated = "" try performMultimodalGeneration( context: context, @@ -614,6 +833,8 @@ import Foundation transcriptEntries: ArraySlice([]) ) } else { + let context = try makeFreshContext(options: runtimeOptions) + defer { llama_free(context) } let maxTokens = structuredOptions.maximumResponseTokens ?? 512 let jsonString = try await generateStructuredJSON( context: context, @@ -664,23 +885,6 @@ import Foundation let runtimeOptions = resolvedOptions(from: options) let maxTokens = runtimeOptions.maximumResponseTokens ?? 100 - let contextParams = createContextParams(from: runtimeOptions) - guard let context = llama_init_from_model(model!, contextParams) else { - throw LlamaLanguageModelError.contextInitializationFailed - } - defer { llama_free(context) } - - // Check if this is an embedding model (no KV cache). - // This early check catches models configured for embeddings that lack a KV cache. - // A complementary architectural check in prepareInitialBatch catches encoder-only - // models (like BERT) by their architecture type. - if llama_get_memory(context) == nil { - throw LlamaLanguageModelError.encoderOnlyModel - } - - // Stabilize runtime behavior per-context - llama_set_causal_attn(context, true) - llama_set_n_threads(context, runtimeOptions.threads, runtimeOptions.threads) var accumulatedText = "" var promptImages: [Data] = [] @@ -705,15 +909,17 @@ import Foundation } if promptImages.isEmpty { - try self.performTextGeneration( - context: context, - model: model!, + try self.generateChatText( + session: session, prompt: fullPrompt, maxTokens: maxTokens, options: runtimeOptions, onToken: yieldToken ) } else { + self.discardCachedSessionContext() + let context = try self.makeFreshContext(options: runtimeOptions) + defer { llama_free(context) } try self.performMultimodalGeneration( context: context, prompt: fullPrompt, @@ -752,6 +958,7 @@ import Foundation llama_backend_init() // Free any existing model before loading a new one + discardCachedSessionContext() if let existingContext = mtmdContext { mtmd_free(existingContext) self.mtmdContext = nil @@ -897,110 +1104,6 @@ import Foundation llama_sampler_chain_add(sampler, llama_sampler_init_dist(options.seed)) } - private func generateText( - context: OpaquePointer, - model: OpaquePointer, - prompt: String, - maxTokens: Int, - options: ResolvedGenerationOptions - ) async throws - -> String - { - guard let vocab = llama_model_get_vocab(model) else { - throw LlamaLanguageModelError.contextInitializationFailed - } - - // Tokenize the prompt - let promptTokens = try tokenizeText(vocab: vocab, text: prompt) - guard !promptTokens.isEmpty else { - throw LlamaLanguageModelError.tokenizationFailed - } - - var batch = llama_batch_init(Int32(options.batchSize), 0, 1) - defer { llama_batch_free(batch) } - - let hasEncoder = try prepareInitialBatch( - batch: &batch, - promptTokens: promptTokens, - model: model, - vocab: vocab, - context: context, - batchSize: options.batchSize, - contextSize: options.contextSize - ) - - // Initialize sampler chain with options - guard let sampler = llama_sampler_chain_init(llama_sampler_chain_default_params()) else { - throw LlamaLanguageModelError.decodingFailed - } - defer { llama_sampler_free(sampler) } - let samplerPtr = UnsafeMutablePointer(sampler) - - let effectiveTemperature = Float(options.temperature) - - // Apply repeat/frequency/presence penalties from custom options - let effectiveRepeatPenalty = options.repeatPenalty - let effectiveRepeatLastN = options.repeatLastN - let effectiveFrequencyPenalty = options.frequencyPenalty - let effectivePresencePenalty = options.presencePenalty - - if effectiveRepeatPenalty != 1.0 || effectiveFrequencyPenalty != 0.0 || effectivePresencePenalty != 0.0 { - llama_sampler_chain_add( - samplerPtr, - llama_sampler_init_penalties( - llama_vocab_n_tokens(vocab), - effectiveRepeatLastN, - effectiveRepeatPenalty, - effectiveFrequencyPenalty, - effectivePresencePenalty - ) - ) - } - - applySampling(sampler: samplerPtr, effectiveTemperature: effectiveTemperature, options: options) - - // Generate tokens one by one - var generatedText = "" - // Track position - for encoder-decoder models, we start from position 1 (after decoder start token) - // For decoder-only models, we continue from the end of the prompt - var n_cur: Int32 = hasEncoder ? 1 : Int32(promptTokens.count) - - for _ in 0 ..< maxTokens { - // Sample next token from logits - llama_batch_get_one creates batch with single token at index 0 - let nextToken = llama_sampler_sample(sampler, context, batch.n_tokens - 1) - llama_sampler_accept(sampler, nextToken) - - // Check for end of sequence - if llama_vocab_is_eog(vocab, nextToken) { - break - } - - // Convert token to text - if let tokenText = tokenToText(vocab: vocab, token: nextToken) { - generatedText += tokenText - } - - // Prepare batch for next token - batch.n_tokens = 1 - batch.token[0] = nextToken - batch.pos[0] = n_cur - batch.n_seq_id[0] = 1 - if let seq_ids = batch.seq_id, let seq_id = seq_ids[0] { - seq_id[0] = 0 - } - batch.logits[0] = 1 - - n_cur += 1 - - let decodeResult = llama_decode(context, batch) - guard decodeResult == 0 else { - break - } - } - - return generatedText - } - /// Builds a JSONSchema-informed prompt for structured output. private func schemaPrompt(for schema: GenerationSchema) -> String { let encoder = JSONEncoder() @@ -1265,23 +1368,21 @@ import Foundation } } - private func performTextGeneration( + private func performTokenGeneration( context: OpaquePointer, - model: OpaquePointer, - prompt: String, + vocab: OpaquePointer, + promptTokens: [llama_token], + startIndex: Int, maxTokens: Int, options: ResolvedGenerationOptions, onToken: (String) -> Void ) throws { - guard let vocab = llama_model_get_vocab(model) else { - throw LlamaLanguageModelError.contextInitializationFailed + guard let model = self.model else { + throw LlamaLanguageModelError.modelLoadFailed } - // Tokenize the prompt - let promptTokens = try tokenizeText(vocab: vocab, text: prompt) - guard !promptTokens.isEmpty else { - throw LlamaLanguageModelError.tokenizationFailed - } + lastReusedTokenCount = startIndex + lastPrefillTokenCount = promptTokens.count - startIndex // Initialize batch var batch = llama_batch_init(Int32(options.batchSize), 0, 1) @@ -1294,7 +1395,8 @@ import Foundation vocab: vocab, context: context, batchSize: options.batchSize, - contextSize: options.contextSize + contextSize: options.contextSize, + startIndex: startIndex ) // Initialize sampler chain with options @@ -1332,6 +1434,7 @@ import Foundation // Track position - for encoder-decoder models, we start from position 1 (after decoder start token) // For decoder-only models, we continue from the end of the prompt var n_cur: Int32 = hasEncoder ? 1 : Int32(promptTokens.count) + var decodedTokens = promptTokens for _ in 0 ..< maxTokens { if Task.isCancelled { @@ -1368,7 +1471,10 @@ import Foundation guard decodeResult == 0 else { break } + decodedTokens.append(nextToken) } + + recordCachedTokens(decodedTokens, context: context) } /// Evaluates a marker-annotated multimodal prompt through the projector, @@ -1542,6 +1648,8 @@ import Foundation /// - context: The model context. /// - batchSize: The batch capacity per decode call. /// - contextSize: The context window the prompt must fit within. + /// - startIndex: The index of the first prompt token to decode. Earlier + /// tokens are already present in the context's state. Defaults to `0`. /// - Returns: `true` if the model has an encoder (for position tracking during generation). /// - Throws: `promptExceedsContextWindow` if the prompt cannot fit in the context window, /// `insufficientMemory` if an encoder prompt exceeds the batch capacity, `encoderOnlyModel` @@ -1553,7 +1661,8 @@ import Foundation vocab: OpaquePointer, context: OpaquePointer, batchSize: UInt32, - contextSize: UInt32 + contextSize: UInt32, + startIndex: Int = 0 ) throws -> Bool { // Leave at least one context cell free for generation. guard promptTokens.count < contextSize else { @@ -1564,7 +1673,7 @@ import Foundation let hasDecoder = llama_model_has_decoder(model) // Encoder models ingest the full prompt in a single llama_encode call. - guard !hasEncoder || promptTokens.count <= batchSize else { + guard !hasEncoder || (startIndex == 0 && promptTokens.count <= batchSize) else { throw LlamaLanguageModelError.insufficientMemory } @@ -1616,7 +1725,7 @@ import Foundation // batch-sized chunks with absolute positions, requesting logits // only for the final token. let capacity = Int(batchSize) - var start = 0 + var start = startIndex while start < promptTokens.count { let count = min(capacity, promptTokens.count - start) batch.n_tokens = Int32(count) diff --git a/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift index ed5a18f6..e1aa4640 100644 --- a/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift @@ -47,6 +47,26 @@ import Testing #expect(!response.content.isEmpty) } + @Test func reusesSessionContextAcrossTurns() async throws { + let session = LanguageModelSession(model: model) + var options = GenerationOptions(maximumResponseTokens: 24) + options[custom: LlamaLanguageModel.self] = .init(contextSize: 2048, batchSize: 512) + + let first = try await session.respond( + to: "My favorite color is blue. Reply with OK.", + options: options + ) + #expect(!first.content.isEmpty) + #expect(model.lastReusedTokenCount == 0) + + let second = try await session.respond( + to: "What is my favorite color? Answer with one word.", + options: options + ) + #expect(!second.content.isEmpty) + #expect(model.lastReusedTokenCount > 0) + } + @Test func customGenerationOptionsRoundTrip() { var options = GenerationOptions( temperature: 0.6,