jkyberneees · Nov 1, 2026 · Nov 1, 2026
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,15 @@
 # Changelog
 
+## 1.1.1
+
+Performance (M2 Pro, measured against 1.1.0 by alternating the two builds, 4 pairs; greedy output is token-identical on Qwen3 0.6B, Qwen2.5 0.5B, Llama 3.2 1B, Phi-3.5 and Mistral 7B):
+
+- The next token is selected on the GPU, inside the same command buffer as the forward pass: repetition penalty, exact top-k, min-p, top-p and a seeded draw. The following decode step is committed before the current one is read back, so the GPU runs step after step without waiting for the CPU to sample, detokenize or launch. Decode at 16 tokens of context: Qwen3 0.6B +10.7%, Llama 3.2 1B +8.5%, Mistral 7B +4.1%; at 1K context +9.0%, +7.8%, +4.0%. Generation as the app and `mlxr run` see it, which previously also paid for CPU sampling: Qwen3 0.6B 238 → 282 tok/s, Llama 3.2 1B 180 → 201 tok/s, Mistral 7B 39.3 → 40.9 tok/s.
+- RMSNorm kernels use vector loads with up to 1024 threads per row; the three norms per layer cost 0.21 ms instead of 0.62 ms on Mistral 7B.
+- The GPU sampler draws from the same distribution as the CPU sampler (tests compare them), but with its own xorshift stream: a given seed produces a different sequence than 1.1.0, and top-k values above 1024 are capped at 1024. Speculative decoding and constrained tool-call names still sample on the CPU.
+- `mlxr bench` reports the pipelined decode, the app's default sampler and the token-selection cost; `profile` adds a per-kernel-class breakdown of one decode step.
+- Prefill is unchanged.
+
 ## 1.1.0
 
 Tool calling in the local server (OpenAI-compatible):

diff --git a/README.md b/README.md
@@ -103,17 +103,17 @@ Each checkpoint above passes end-to-end generation and the prefill-versus-decode
 
 ## Performance
 
-Measured with `mlxr bench` on an **M2 Pro (16 GB)**. Decode is measured with greedy sampling.
+Measured with `mlxr bench` on an **M2 Pro (16 GB)**. Decode is measured with greedy sampling, as generation runs it: the token is selected on the GPU and the next step is queued before the current one is read back, so these are also the rates the app and `mlxr run` reach.
 
 | Model (4-bit) | Weights | Prefill (512 tok) | Decode @16 | Decode @1K ctx |
 |---|---|---|---|---|
-| Qwen2.5 0.5B Instruct | 278 MB | 2,960 tok/s | 307 tok/s | 260 tok/s |
-| Qwen3 0.6B | 336 MB | 2,040 tok/s | 252 tok/s | 214 tok/s |
-| Llama 3.2 1B Instruct | 695 MB | 1,180 tok/s | 187 tok/s | 174 tok/s |
-| Phi-3.5 mini | 2.15 GB | 320 tok/s | 67 tok/s | 59 tok/s |
-| Mistral 7B v0.3 | 4.08 GB | 166 tok/s | 39.4 tok/s | 37.9 tok/s |
+| Qwen2.5 0.5B Instruct | 278 MB | 2,960 tok/s | 358 tok/s | 295 tok/s |
+| Qwen3 0.6B | 336 MB | 2,040 tok/s | 289 tok/s | 239 tok/s |
+| Llama 3.2 1B Instruct | 695 MB | 1,180 tok/s | 203 tok/s | 188 tok/s |
+| Phi-3.5 mini | 2.15 GB | 320 tok/s | 70 tok/s | 61 tok/s |
+| Mistral 7B v0.3 | 4.08 GB | 166 tok/s | 41.0 tok/s | 39.5 tok/s |
 
-**The 7B ceiling.** Decoding one token reads every weight once, so speed is capped by memory bandwidth. This M2 Pro measures 197 GB/s with a plain read kernel, which caps a 4.08 GB model at about 48 tok/s. MLX Runner reaches 39.4 tok/s, about 82% of that cap. Reaching 50+ tok/s on a 4-bit 7B model therefore needs one of: a Mac with more bandwidth (M-series Max chips have about 400 GB/s), a 3-bit model (which would raise the ceiling to about 60 tok/s, but only 4-bit has an optimized decode kernel so far), or speculative decoding (see Limitations).
+**The 7B ceiling.** Decoding one token reads every weight once, so speed is capped by memory bandwidth. This M2 Pro measures 197 GB/s with a plain read kernel, which caps a 4.08 GB model at about 48 tok/s. MLX Runner reaches 41.0 tok/s, about 85% of that cap. Reaching 50+ tok/s on a 4-bit 7B model therefore needs one of: a Mac with more bandwidth (M-series Max chips have about 400 GB/s), a 3-bit model (which would raise the ceiling to about 60 tok/s, but only 4-bit has an optimized decode kernel so far), or speculative decoding (see Limitations).
 
 Load times stay short because the weights are `mmap`ed and wrapped in Metal buffers with no copy: about 0.5 s for the small models, 1.6 s for Phi-3.5 and 3.3 s for Mistral 7B. Those figures include a warm-up pass that pages the weights in.
 
@@ -154,7 +154,7 @@ Sources/
     Transformer.swift   Model graph: one command buffer per step, fused kernels
     Tokenizer.swift     BPE from tokenizer.json (byte-level and SentencePiece)
     ChatTemplate.swift  ChatML / Llama 3 / Mistral / Phi-3 prompt formats
-    Engine.swift        Generation loop, sampler, KV-prefix reuse, cancellation, context fitting
+    Engine.swift        Generation loop, pipelined decode, KV-prefix reuse, cancellation, context fitting
     Hub.swift           Hugging Face API client and downloader
     ModelStore.swift    Local library and HF-cache scanning
     OpenAIServer.swift  OpenAI-compatible HTTP server (Network.framework)
@@ -170,6 +170,7 @@ How the engine works:
   - Per-head RMSNorm, RoPE and the KV-cache append run as a single kernel.
   - The residual add is fused with the next RMSNorm.
   - Attention uses online softmax across 16 to 32 simdgroups.
+  - The token is selected on the GPU at the end of the same command buffer (repetition penalty, exact top-k, min-p, top-p, one seeded draw) and written to a token-history buffer that the next step's embedding lookup reads. The next step is therefore committed before the current one is read back, and the GPU never waits for sampling or detokenization.
 - **Prefill:**
   - Quantized GEMMs on `simdgroup_matrix` tiles; each weight tile is dequantized once and shared by 32 tokens.
   - Flash attention on simdgroup matrices: 8×8 tensor ops for Q·Kᵀ and P·V, with the softmax rescale done as a diagonal-matrix multiply.
@@ -181,6 +182,7 @@ How the engine works:
 
 - GPU quantized matmul kernels for every bit width, group size and scale dtype, across all three code paths (decode matvec, small-batch, prefill GEMM), each compared against a CPU reference.
 - Fused multi-matrix decode with bias, dense f16 matmul, and the quantized embedding lookup.
+- GPU token selection against the CPU sampler: greedy (ties, repetition penalty), the drawn distribution for several top-k / top-p / min-p settings, ties at the top, tiny vocabularies and non-finite logits.
 - Tokenizer known ids and Unicode round-trips; chat templates; sampler behavior (greedy, repetition penalty, top-k, top-p, seeding); the safetensors loader, including misaligned tensors.
 - On real models:
   - Prefill and token-by-token decode produce the same next-token distribution, for all five architectures.

diff --git a/Sources/MLXEngine/Engine.swift b/Sources/MLXEngine/Engine.swift
@@ -400,11 +400,33 @@ public final class InferenceEngine: @unchecked Sendable {
         while common < limit && cachedTokens[common] == tokens[common] { common += 1 }
         stats.cachedPromptTokens = common
 
+        let vocab = model.vocabSize
+        let count = min(vocab, tok.vocabularySize)
+        let seed = parameters.seed ?? UInt64.random(in: 1...UInt64.max)
+        // Constrained name: the text generated so far while a forced tool-call name is open.
+        var nameText = ""
+        var nameOpen = !(nameChoices ?? []).isEmpty
+        let texts = nameOpen ? tokenTextIndex(tok) : nil
+        // Speculative decoding (prompt lookup): verifying k drafted tokens costs one batched
+        // pass, so every accepted draft is a token generated without its own weight read.
+        let speculative = (parameters.speculative ?? false) && model.supportsBatchedVerify && !nameOpen
+        // Plain decoding selects tokens on the GPU, so each step is queued before the previous one
+        // is read back and the GPU never waits for sampling or detokenisation.
+        let gpuSampling = !speculative && !nameOpen
+        let gpuParams = SampleParams(parameters, vocab: count)
+
         cachedTokens = Array(tokens[0..<common])
         let t0 = Date()
-        var logits: UnsafeMutablePointer<Float>
+        var logits: UnsafeMutablePointer<Float>?   // CPU sampling only
+        var firstToken: Int?
         do {
-            logits = try model.forward(tokens[common...].map { Int32($0) }, offset: common, isCancelled: { token.isCancelled })
+            let prompt = tokens[common...].map { Int32($0) }
+            if gpuSampling {
+                model.seedSampler(seed)
+                firstToken = try model.prefill(prompt, offset: common, sample: gpuParams, isCancelled: { token.isCancelled })
+            } else {
+                logits = try model.forward(prompt, offset: common, isCancelled: { token.isCancelled })
+            }
         } catch EngineError.cancelled {
             stats.stopReason = "cancelled"
             emit(.finished(stats))
@@ -417,22 +439,13 @@ public final class InferenceEngine: @unchecked Sendable {
         stats.promptTokensPerSecond = Double(tokens.count - common) / max(tPrompt, 1e-6)
         emit(.promptProcessed(stats))
 
-        var sampler = Sampler(params: parameters, rng: XorShift(seed: parameters.seed ?? UInt64.random(in: 1...UInt64.max)))
+        var sampler = Sampler(params: parameters, rng: XorShift(seed: seed))
         var decoder = StreamingDecoder()
         let tGen = Date()
         stats.stopReason = "length"
         let maxNew = min(parameters.maxTokens, model.maxSeq - tokens.count)
-        let vocab = model.vocabSize
-        let count = min(vocab, tok.vocabularySize)
         var atStart = true
         var drafted = 0, acceptedDrafts = 0
-        // Speculative decoding (prompt lookup): verifying k drafted tokens costs one batched
-        // pass, so every accepted draft is a token generated without its own weight read.
-        // Constrained name: the text generated so far while a forced tool-call name is open.
-        var nameText = ""
-        var nameOpen = !(nameChoices ?? []).isEmpty
-        let texts = nameOpen ? tokenTextIndex(tok) : nil
-        let speculative = (parameters.speculative ?? false) && model.supportsBatchedVerify && !nameOpen
         var draftLen = 4
 
         /// Emits a sampled token. Returns false when generation must stop.
@@ -467,19 +480,44 @@ public final class InferenceEngine: @unchecked Sendable {
             return d
         }
 
-        // Invariant: `history.last` (the pending token) is not yet in the KV cache;
-        // everything before it is, and `cachedTokens` mirrors exactly that.
         var running = maxNew > 0 && !token.isCancelled
-        if running { running = accept(sampler.draw(try dist(logits))) }
+        if let first = firstToken {
+            // GPU sampling: the step for the newest token is committed before the previous step is
+            // read back, so the GPU runs them back to back. A step that turns out to process a stop
+            // token only writes a K/V row that `cachedTokens` does not claim.
+            if running { running = accept(first) }
+            var cachedCount = tokens.count   // positions whose K/V rows the cache holds
+            var inflight: Transformer.DecodeStep? = running ? try model.beginDecodeStep(offset: history.count - 1, sample: gpuParams) : nil
+            while let cur = inflight {
+                inflight = nil
+                let nextOffset = cur.offset + 1
+                if !token.isCancelled, stats.generatedTokens + 1 < maxNew, nextOffset + 1 < model.maxSeq {
+                    do { inflight = try model.beginDecodeStep(offset: nextOffset, sample: gpuParams) }
+                    catch { _ = try? model.finish(cur); throw error }
+                }
+                let t: Int
+                do { t = try model.finish(cur) } catch { if let f = inflight { _ = try? model.finish(f) }; throw error }
+                cachedCount = cur.offset + 1
+                if token.isCancelled { stats.stopReason = "cancelled"; break }
+                if !accept(t) { break }
+            }
+            if let f = inflight, (try? model.finish(f)) != nil { cachedCount = f.offset + 1 }
+            cachedTokens = Array(history.prefix(min(history.count, cachedCount)))
+            running = false
+        }
+        // CPU sampling (speculative decoding, constrained tool names). Invariant: `history.last`
+        // (the pending token) is not yet in the KV cache; everything before it is, and
+        // `cachedTokens` mirrors exactly that.
+        if running, let logits { running = accept(sampler.draw(try dist(logits))) }
         while running {
             if token.isCancelled { stats.stopReason = "cancelled"; break }
             let pending = history[history.count - 1]
             let room = min(maxNew - stats.generatedTokens, model.maxSeq - history.count) - 1
             let draft = speculative ? NGramDraft.propose(history, maxTokens: min(draftLen, room, Transformer.maxVerifyTokens - 1)) : []
             if draft.isEmpty {
-                logits = try model.forward([Int32(pending)], offset: history.count - 1)
+                let row = try model.forward([Int32(pending)], offset: history.count - 1)
                 cachedTokens = history
-                running = accept(sampler.draw(try dist(logits)))
+                running = accept(sampler.draw(try dist(row)))
                 continue
             }
             let rows = try model.forward(([pending] + draft).map { Int32($0) }, offset: history.count - 1, allLogits: true)

diff --git a/Sources/MLXEngine/GPU.swift b/Sources/MLXEngine/GPU.swift
@@ -38,6 +38,19 @@ struct QKVPostParams { var D: Int32; var nH: Int32; var nKV: Int32; var rowStrid
 struct NormParams { var dim: Int32; var headsPerRow: Int32; var rowStride: Int32; var eps: Float }
 struct AttnParams { var D: Int32; var nH: Int32; var nKV: Int32; var qStride: Int32; var maxSeq: Int32; var offset: Int32; var scale: Float; var outStride: Int32; var T: Int32 = 1 }
 struct EWParams { var n: Int32; var stride: Int32; var cols: Int32; var pad: Int32 = 0 }
+/// GPU token selection (`sample_token`); `V` is the number of logits considered, `pos` the
+/// position of the token whose logits these are (the result goes to hist[pos + 1]).
+public struct SampleParams {
+    var V: Int32; var pos: Int32; var window: Int32; var topK: Int32
+    var penalty: Float; var temperature: Float; var topP: Float; var minP: Float
+
+    /// `vocab`: how many logits to consider (the tokenizer's vocabulary may be smaller than lm_head).
+    public init(_ p: SamplingParameters, vocab: Int) {
+        V = Int32(vocab); pos = 0
+        window = Int32(max(0, p.repetitionWindow)); topK = Int32(p.topK)
+        penalty = p.repetitionPenalty; temperature = p.temperature; topP = p.topP; minP = p.minP
+    }
+}
 
 /// Shared Metal state: device, queue, and the compiled kernel library.
 public final class GPU: @unchecked Sendable {