Skip to content

Ingest decoder prompts in batch-sized chunks #196

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

Open
james-333i wants to merge 2 commits into
base: main
Choose a base branch
from
+82 −30
Open

Ingest decoder prompts in batch-sized chunks #196

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3e168e3
Fix LlamaLanguageModel build against current llama.swift
james-333i Aug 25, 2026
e375911
Ingest decoder prompts in batch-sized chunks
james-333i Aug 25, 2026
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
9 changes: 9 additions & 0 deletions Package.resolved
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

87 changes: 57 additions & 30 deletions Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down Expand Up @@ -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
Expand All @@ -829,6 +829,7 @@ import Foundation
llama_sampler_chain_add(
samplerPtr,
llama_sampler_init_penalties(
llama_vocab_n_tokens(vocab),
effectiveRepeatLastN,
effectiveRepeatPenalty,
effectiveFrequencyPenalty,
Expand All @@ -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
Expand Down Expand Up @@ -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 {
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -1197,6 +1201,7 @@ import Foundation
llama_sampler_chain_add(
samplerPtr,
llama_sampler_init_penalties(
llama_vocab_n_tokens(vocab),
effectiveRepeatLastN,
effectiveRepeatPenalty,
effectiveFrequencyPenalty,
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
}
}

Expand Down Expand Up @@ -1537,6 +1561,7 @@ import Foundation
case decodingFailed
case invalidModelPath
case insufficientMemory
case promptExceedsContextWindow
case unsupportedFeature
case encoderOnlyModel

Expand All @@ -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:
Expand Down
16 changes: 16 additions & 0 deletions Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down