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..efcb5222 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 } @@ -807,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 @@ -829,6 +829,7 @@ import Foundation llama_sampler_chain_add( samplerPtr, llama_sampler_init_penalties( + llama_vocab_n_tokens(vocab), effectiveRepeatLastN, effectiveRepeatPenalty, effectiveFrequencyPenalty, @@ -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 { @@ -958,6 +960,7 @@ import Foundation llama_sampler_chain_add( samplerPointer, llama_sampler_init_penalties( + llama_vocab_n_tokens(vocab), options.repeatLastN, options.repeatPenalty, options.frequencyPenalty, @@ -968,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, @@ -1174,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 @@ -1197,6 +1201,7 @@ import Foundation llama_sampler_chain_add( samplerPtr, llama_sampler_init_penalties( + llama_vocab_n_tokens(vocab), effectiveRepeatLastN, effectiveRepeatPenalty, effectiveFrequencyPenalty, @@ -1211,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 @@ -1272,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) @@ -1341,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 } } @@ -1537,6 +1561,7 @@ import Foundation case decodingFailed case invalidModelPath case insufficientMemory + case promptExceedsContextWindow case unsupportedFeature case encoderOnlyModel @@ -1556,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,