From 3e168e39c4ae968175fcea546b057cc0df5a7799 Mon Sep 17 00:00:00 2001 From: james-333i Date: Tue, 25 Aug 2026 11:29:46 -0700 Subject: [PATCH 1/2] Fix LlamaLanguageModel build against current llama.swift The open-ended dependency range resolves llama.swift to releases wrapping current llama.cpp builds, where the Llama trait no longer compiles: llama_sampler_init_penalties regained its leading n_vocab parameter, and llama_model_params replaced use_mmap and use_mlock with a llama_load_mode enum. Pass the vocabulary size at all three penalties call sites and set load_mode to LLAMA_LOAD_MODE_MMAP, matching the previous mmap-only behavior. Verified against llama.swift 2.10549.0 with the full live test suite. --- Package.resolved | 9 +++++++++ Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift | 6 ++++-- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/Package.resolved b/Package.resolved index b689a691..0b86bc46 100644 --- a/Package.resolved +++ b/Package.resolved @@ -19,6 +19,15 @@ "version" : "1.3.1" } }, + { + "identity" : "llama.swift", + "kind" : "remoteSourceControl", + "location" : "https://github.com/mattt/llama.swift", + "state" : { + "revision" : "716419d4d7aa542fce301e809cde7234c68ddbc6", + "version" : "2.10549.0" + } + }, { "identity" : "partialjsondecoder", "kind" : "remoteSourceControl", diff --git a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift index 6bcd8045..56d858a9 100644 --- a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift @@ -673,8 +673,7 @@ import Foundation params.n_gpu_layers = 0 // Try to reduce memory usage - params.use_mmap = true - params.use_mlock = false + params.load_mode = LLAMA_LOAD_MODE_MMAP return params } @@ -829,6 +828,7 @@ import Foundation llama_sampler_chain_add( samplerPtr, llama_sampler_init_penalties( + llama_vocab_n_tokens(vocab), effectiveRepeatLastN, effectiveRepeatPenalty, effectiveFrequencyPenalty, @@ -958,6 +958,7 @@ import Foundation llama_sampler_chain_add( samplerPointer, llama_sampler_init_penalties( + llama_vocab_n_tokens(vocab), options.repeatLastN, options.repeatPenalty, options.frequencyPenalty, @@ -1197,6 +1198,7 @@ import Foundation llama_sampler_chain_add( samplerPtr, llama_sampler_init_penalties( + llama_vocab_n_tokens(vocab), effectiveRepeatLastN, effectiveRepeatPenalty, effectiveFrequencyPenalty, From e37591106253e231e422b886100d38ab2e65ffef Mon Sep 17 00:00:00 2001 From: james-333i Date: Tue, 25 Aug 2026 12:00:39 -0700 Subject: [PATCH 2/2] Ingest decoder prompts in batch-sized chunks Prompts longer than the batch capacity (512 tokens by default) threw insufficientMemory before generation started, so multi-turn conversations failed as soon as the rendered chat history crossed the batch size, regardless of how much memory was actually available. Feed decoder-only prompts through llama_decode in batch-sized chunks with absolute positions, requesting logits only for the final token. Generation positions now derive from the full prompt length rather than the last batch's token count. Encoder models keep the single-batch requirement, and a prompt that cannot fit in the context window now fails with a new promptExceedsContextWindow error instead of a misleading memory error. Adds a live test generating from a prompt several times the batch size. --- .../Models/LlamaLanguageModel.swift | 81 ++++++++++++------- .../LlamaLanguageModelTests.swift | 16 ++++ 2 files changed, 69 insertions(+), 28 deletions(-) diff --git a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift index 56d858a9..efcb5222 100644 --- a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift @@ -806,7 +806,8 @@ import Foundation model: model, vocab: vocab, context: context, - batchSize: options.batchSize + batchSize: options.batchSize, + contextSize: options.contextSize ) // Initialize sampler chain with options @@ -843,7 +844,7 @@ import Foundation 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 : batch.n_tokens + 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 @@ -945,7 +946,8 @@ import Foundation model: model!, vocab: vocab, context: context, - batchSize: options.batchSize + batchSize: options.batchSize, + contextSize: options.contextSize ) guard let sampler = llama_sampler_chain_init(llama_sampler_chain_default_params()) else { @@ -969,7 +971,7 @@ import Foundation applySampling(sampler: samplerPointer, effectiveTemperature: options.temperature, options: options) let vocabSize = Int(llama_vocab_n_tokens(vocab)) - let initialPosition: Int32 = hasEncoder ? 1 : batchPointer.pointee.n_tokens + let initialPosition: Int32 = hasEncoder ? 1 : Int32(promptTokens.count) let backend = LlamaTokenBackend( context: context, @@ -1175,7 +1177,8 @@ import Foundation model: model, vocab: vocab, context: context, - batchSize: options.batchSize + batchSize: options.batchSize, + contextSize: options.contextSize ) // Initialize sampler chain with options @@ -1213,7 +1216,7 @@ import Foundation // Generate tokens one by one // 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 : batch.n_tokens + var n_cur: Int32 = hasEncoder ? 1 : Int32(promptTokens.count) for _ in 0 ..< maxTokens { // Sample next token from logits of the last token we just decoded @@ -1274,31 +1277,43 @@ import Foundation /// Prepares the initial batch for text generation, handling encoder-decoder vs decoder-only models. /// + /// Decoder-only prompts longer than the batch capacity are ingested in + /// batch-sized chunks. Encoder models must fit the prompt in one batch. + /// /// - Parameters: /// - batch: The batch to prepare (must be initialized with sufficient capacity). /// - promptTokens: The tokenized prompt tokens. /// - model: The loaded model. /// - vocab: The model vocabulary. /// - context: The model context. - /// - batchSize: The batch capacity to validate against (prevents buffer overflow). + /// - batchSize: The batch capacity per decode call. + /// - contextSize: The context window the prompt must fit within. /// - Returns: `true` if the model has an encoder (for position tracking during generation). - /// - Throws: `insufficientMemory` if prompt token count exceeds batch capacity, `encoderOnlyModel` if the model cannot generate text, `encodingFailed` or `decodingFailed` on failure. + /// - Throws: `promptExceedsContextWindow` if the prompt cannot fit in the context window, + /// `insufficientMemory` if an encoder prompt exceeds the batch capacity, `encoderOnlyModel` + /// if the model cannot generate text, `encodingFailed` or `decodingFailed` on failure. private func prepareInitialBatch( batch: inout llama_batch, promptTokens: [llama_token], model: OpaquePointer, vocab: OpaquePointer, context: OpaquePointer, - batchSize: UInt32 + batchSize: UInt32, + contextSize: UInt32 ) throws -> Bool { - // Validate that prompt token count doesn't exceed batch capacity to prevent buffer overflow - guard promptTokens.count <= batchSize else { - throw LlamaLanguageModelError.insufficientMemory + // Leave at least one context cell free for generation. + guard promptTokens.count < contextSize else { + throw LlamaLanguageModelError.promptExceedsContextWindow } let hasEncoder = llama_model_has_encoder(model) 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 { + throw LlamaLanguageModelError.insufficientMemory + } + if hasEncoder { // For encoder models, first encode the prompt batch.n_tokens = Int32(promptTokens.count) @@ -1343,25 +1358,32 @@ import Foundation throw LlamaLanguageModelError.encoderOnlyModel } } else { - // Standard decoder-only model (most LLMs) - batch.n_tokens = Int32(promptTokens.count) - for i in 0 ..< promptTokens.count { - let idx = Int(i) - batch.token[idx] = promptTokens[idx] - batch.pos[idx] = Int32(i) - batch.n_seq_id[idx] = 1 - if let seq_ids = batch.seq_id, let seq_id = seq_ids[idx] { - seq_id[0] = 0 + // Standard decoder-only model (most LLMs): feed the prompt in + // batch-sized chunks with absolute positions, requesting logits + // only for the final token. + let capacity = Int(batchSize) + var start = 0 + while start < promptTokens.count { + let count = min(capacity, promptTokens.count - start) + batch.n_tokens = Int32(count) + for i in 0 ..< count { + batch.token[i] = promptTokens[start + i] + batch.pos[i] = Int32(start + i) + batch.n_seq_id[i] = 1 + if let seq_ids = batch.seq_id, let seq_id = seq_ids[i] { + seq_id[0] = 0 + } + batch.logits[i] = 0 } - batch.logits[idx] = 0 - } - if batch.n_tokens > 0 { - batch.logits[Int(batch.n_tokens) - 1] = 1 - } + if start + count == promptTokens.count { + batch.logits[count - 1] = 1 + } - guard llama_decode(context, batch) == 0 else { - throw LlamaLanguageModelError.decodingFailed + guard llama_decode(context, batch) == 0 else { + throw LlamaLanguageModelError.decodingFailed + } + start += count } } @@ -1539,6 +1561,7 @@ import Foundation case decodingFailed case invalidModelPath case insufficientMemory + case promptExceedsContextWindow case unsupportedFeature case encoderOnlyModel @@ -1558,6 +1581,8 @@ import Foundation return "Invalid model file path" case .insufficientMemory: return "Insufficient memory for operation" + case .promptExceedsContextWindow: + return "Prompt is longer than the model's context window" case .unsupportedFeature: return "This LlamaLanguageModel does not support image segments" case .encoderOnlyModel: diff --git a/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift index 3d9e32c7..8c6f947d 100644 --- a/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift @@ -27,6 +27,22 @@ import Testing #expect(customModel.repeatLastN == 64) } + @Test func promptLongerThanBatchSize() async throws { + let session = LanguageModelSession(model: model) + var options = GenerationOptions(maximumResponseTokens: 16) + options[custom: LlamaLanguageModel.self] = .init(batchSize: 32) + + let filler = Array( + repeating: "The quick brown fox jumps over the lazy dog.", + count: 30 + ).joined(separator: " ") + let response = try await session.respond( + to: "\(filler)\n\nReply with a single word.", + options: options + ) + #expect(!response.content.isEmpty) + } + @Test func customGenerationOptionsRoundTrip() { var options = GenerationOptions( temperature: 0.6,