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..03fbd609 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,
@@ -1109,18 +1110,37 @@ import Foundation
             }
 
             mutating func sample(from allowedTokens: Set<Int>) async throws -> Int {
-                guard let logits = llama_get_logits(context) else {
+                // Masking llama_get_logits in place does not constrain
+                // llama_sampler_sample: the context can refresh that buffer when
+                // the sampler fetches logits, dropping the mask. Build a candidate
+                // array holding only the allowed tokens and run the sampler chain
+                // over it instead.
+                guard let logits = llama_get_logits_ith(context, batch.pointee.n_tokens - 1) else {
                     return eosToken
                 }
 
-                for tokenIndex in 0 ..< vocabSize {
-                    if !allowedTokens.contains(tokenIndex) {
-                        logits[tokenIndex] = -Float.infinity
-                    }
+                var candidates = allowedTokens.compactMap { token -> llama_token_data? in
+                    guard token >= 0, token < vocabSize else { return nil }
+                    return llama_token_data(id: llama_token(token), logit: logits[token], p: 0)
+                }
+                guard !candidates.isEmpty else {
+                    return eosToken
                 }
 
-                let tokenIndex = batch.pointee.n_tokens - 1
-                return Int(llama_sampler_sample(sampler, context, tokenIndex))
+                let selectedToken = candidates.withUnsafeMutableBufferPointer { buffer -> Int in
+                    var array = llama_token_data_array(
+                        data: buffer.baseAddress,
+                        size: buffer.count,
+                        selected: -1,
+                        sorted: false
+                    )
+                    llama_sampler_apply(sampler, &array)
+                    guard array.selected >= 0, array.selected < Int64(array.size), let data = array.data else {
+                        return eosToken
+                    }
+                    return Int(data[Int(array.selected)].id)
+                }
+                return selectedToken
             }
         }
 
@@ -1197,6 +1217,7 @@ import Foundation
                     llama_sampler_chain_add(
                         samplerPtr,
                         llama_sampler_init_penalties(
+                            llama_vocab_n_tokens(vocab),
                             effectiveRepeatLastN,
                             effectiveRepeatPenalty,
                             effectiveFrequencyPenalty,