Everything Bagel

Sources/FFAI/CacheSnapshot.swift

@@ -0,0 +1,117 @@
+// CacheSnapshot — composite snapshot/restore over `[any LayerCacheProtocol]`.
+//
+// Speculative decode needs to roll back ALL of a model's per-layer
+// caches on draft rejection: KVCache (attention), GDNStateCache +
+// ConvStateCache (GDN). This file gives the spec-decode driver one
+// call to snapshot every layer + one call to restore.
+//
+// Per-cache-kind snapshot payload:
+//   * KVCache              → `length` integer only (the KV slots
+//                            beyond `length` are unused; appending
+//                            after restore overwrites them).
+//   * GDNStateCache        → Tensor snapshot of `current` +
+//                            `length` integer.
+//   * ConvStateCache       → Tensor snapshot of `state`.
+//   * Qwen35GDNLayerCache  → both of the above (conv + GDN +
+//                            shared `length`).
+//
+// Cost on Qwen3.6-A3B: 60 MiB GDN state snapshots + 900 KiB conv
+// snapshots + ~zero attention metadata = ~61 MiB per spec step. On
+// Apple silicon unified memory at ~400 GB/s that's ~150 µs in
+// aggregate blit time — negligible vs ~60 ms decode step.
+
+import Foundation
+
+/// Per-layer snapshot. The enum cases match the concrete cache types
+/// the model layer slots take; the spec-decode driver routes per
+/// layer.
+public enum LayerCacheSnapshot {
+    case kv(length: Int, absolutePosition: Int)
+    case gdn(currentState: Tensor, length: Int)
+    case conv(state: Tensor)
+    case gdnLayer(conv: Tensor, gdnState: Tensor, length: Int)
+}
+
+public extension Array where Element == any LayerCacheProtocol {
+    /// Snapshot every layer cache. Cost is ~150 µs aggregate blit time
+    /// on Qwen3.6-A3B (60 MiB GDN + 900 KB conv).
+    func snapshotAll(device: Device = .shared) -> [LayerCacheSnapshot] {
+        return self.map { cache -> LayerCacheSnapshot in
+            if let composite = cache as? Qwen35GDNLayerCache {
+                return .gdnLayer(
+                    conv: composite.conv.snapshot(device: device),
+                    gdnState: composite.gdn.snapshot(device: device),
+                    length: composite.length)
+            }
+            if let gdn = cache as? GDNStateCache {
+                return .gdn(
+                    currentState: gdn.snapshot(device: device),
+                    length: gdn.length)
+            }
+            if let conv = cache as? ConvStateCache {
+                return .conv(state: conv.snapshot(device: device))
+            }
+            if let kv = cache as? KVCache {
+                return .kv(length: kv.length, absolutePosition: kv.absolutePosition)
+            }
+            preconditionFailure(
+                "CacheSnapshot: unhandled LayerCacheProtocol subtype \(type(of: cache)). Add a case if a new cache kind ships."
+            )
+        }
+    }
+
+    /// Restore every layer cache from a snapshot taken via `snapshotAll`.
+    /// Caller must pass the same `[any LayerCacheProtocol]` instance
+    /// the snapshot came from — order + identity must match.
+    func restoreAll(from snapshots: [LayerCacheSnapshot],
+                    device: Device = .shared) {
+        precondition(self.count == snapshots.count,
+                     "CacheSnapshot.restoreAll: layer count mismatch (caches=\(self.count), snapshots=\(snapshots.count))")
+        for (cache, snap) in zip(self, snapshots) {
+            switch (cache, snap) {
+            case let (composite as Qwen35GDNLayerCache, .gdnLayer(convT, gdnT, length)):
+                composite.conv.restore(from: convT, device: device)
+                composite.gdn.restore(from: gdnT, device: device)
+                // Truncate `length` back to the snapshotted value.
+                // Qwen35GDNLayerCache only exposes `advance()`; rewind
+                // by re-creating the counter via internal state — see
+                // the type's `restoreLength` accessor below.
+                composite.restoreLength(to: length)
+            case let (gdn as GDNStateCache, .gdn(currentT, length)):
+                gdn.restore(from: currentT, device: device)
+                gdn.restoreLength(to: length)
+            case let (conv as ConvStateCache, .conv(stateT)):
+                conv.restore(from: stateT, device: device)
+            case let (kv as KVCache, .kv(length, _)):
+                kv.truncate(toLength: length)
+            default:
+                preconditionFailure(
+                    "CacheSnapshot.restoreAll: cache/snapshot mismatch at layer — cache=\(type(of: cache)), snapshot=\(snap)"
+                )
+            }
+        }
+    }
+}
+
+// `length` rewind helpers — `swap()` / `advance()` only move forward.
+// Spec decode needs to move backward on reject.
+
+public extension GDNStateCache {
+    /// Rewind the position counter without touching tensor state.
+    /// Used by spec-decode AFTER the buffers have been restored via
+    /// `restore(from:)`. Was previously implemented as reset+swap loop,
+    /// which zeroed the state buffers — see `setLength` for the safe
+    /// path.
+    func restoreLength(to length: Int) {
+        self.setLength(length)
+    }
+}
+
+public extension Qwen35GDNLayerCache {
+    /// Composite-level length-only restore. Tensor state is restored
+    /// separately by the caller; this fixes the position counter
+    /// without disturbing the just-restored conv/gdn buffers.
+    func restoreLength(to length: Int) {
+        self.setLength(length)
+    }
+}

Sources/FFAI/ConvStateCache.swift

@@ -30,11 +30,71 @@ public final class ConvStateCache: @unchecked Sendable {
         self.state = Tensor.empty(shape: [kernelSize - 1, nChannels],
                                   dtype: dtype, device: device)
         self.state.zero()
+        // Conv state lives for the lifetime of a generation. Pin it in
+        // the device residency set — conv1dCausalStep touches it every
+        // token, every layer.
+        device.markWeightsResident([self.state.buffer])
     }
 
     /// Reset to zero. Used between sessions; cheap (small fp16/bf16 buffer).
     public func reset() { state.zero() }
 
+    /// Snapshot the current rolling window to a fresh `Tensor`. Used by
+    /// speculative decode — the conv1d_causal_step kernel mutates this
+    /// window in-place each token (drops state[0], appends current
+    /// input as state[K-2]), and the mutation isn't trivially
+    /// reversible without remembering the exact dropped slot and old
+    /// state contents. Snapshot/restore is the cleanest rollback.
+    ///
+    /// Cost: `(kernelSize - 1) * nChannels` × dtype-size bytes per
+    /// layer. At Qwen3.6-A3B convDim=5120, kernelSize=4, bf16:
+    /// 3 * 5120 * 2 = 30 720 bytes per layer × 30 GDN layers ≈ 900 KB
+    /// — small enough to snapshot every spec step without measurable
+    /// host overhead.
+    /// Cached snapshot tensor (ITER 33 — bound spec-decode churn).
+    private var cachedSnapshot: Tensor?
+
+    public func snapshot(device: Device = .shared) -> Tensor {
+        if cachedSnapshot == nil {
+            cachedSnapshot = Tensor.empty(shape: [kernelSize - 1, nChannels],
+                                           dtype: dtype, device: device)
+        }
+        let snap = cachedSnapshot!
+        let cmd = device.makeCommandBuffer()
+        guard let blit = cmd.makeBlitCommandEncoder() else {
+            preconditionFailure("ConvStateCache.snapshot: makeBlitCommandEncoder failed")
+        }
+        let bytes = (kernelSize - 1) * nChannels * dtype.byteSize
+        blit.copy(from: state.buffer, sourceOffset: state.offset,
+                  to: snap.buffer, destinationOffset: snap.offset,
+                  size: bytes)
+        blit.endEncoding()
+        cmd.commit()
+        cmd.waitUntilCompleted()
+        return snap
+    }
+
+    /// Restore the rolling window from a snapshot. Overwrites the
+    /// state buffer with the snapshot contents.
+    public func restore(from snapshot: Tensor, device: Device = .shared) {
+        let expected = (kernelSize - 1) * nChannels
+        precondition(snapshot.elementCount == expected,
+                     "ConvStateCache.restore: snapshot has \(snapshot.elementCount) elements, expected \(expected)")
+        precondition(snapshot.dtype == dtype,
+                     "ConvStateCache.restore: snapshot dtype \(snapshot.dtype) ≠ state dtype \(dtype)")
+        let cmd = device.makeCommandBuffer()
+        guard let blit = cmd.makeBlitCommandEncoder() else {
+            preconditionFailure("ConvStateCache.restore: makeBlitCommandEncoder failed")
+        }
+        let bytes = expected * dtype.byteSize
+        blit.copy(from: snapshot.buffer, sourceOffset: snapshot.offset,
+                  to: state.buffer, destinationOffset: state.offset,
+                  size: bytes)
+        blit.endEncoding()
+        cmd.commit()
+        cmd.waitUntilCompleted()
+    }
+
     /// Bytes occupied by the rolling window. Constant w.r.t. sequence
     /// length — that's the streaming-decode design.
     public var bytesAllocated: Int {

Sources/FFAI/Device.swift

@@ -10,6 +10,15 @@ public final class Device: @unchecked Sendable {
     public let mtlDevice: MTLDevice
     public let commandQueue: MTLCommandQueue
 
+    /// Lazy MTLResidencySet holding the model's weight buffers. Marked
+    /// resident after `Model.load` finishes; saves per-command-buffer
+    /// residency tracking on every prefill / decode dispatch. Disable
+    /// via `FFAI_NO_RESIDENCY_SET=1`. Stored as `Any?` so the deployment
+    /// target stays < macOS 15; the actual cast lives inside the
+    /// `@available` block below.
+    private var weightResidencySet: Any?
+    private let residencyLock = NSLock()
+
     public static let shared: Device = {
         // Reuse the same MTLDevice + queue MetalTileSwift uses, so PSOs
         // and buffers are guaranteed compatible.
@@ -37,4 +46,38 @@ public final class Device: @unchecked Sendable {
         }
         return cb
     }
+
+    /// Mark `buffers` as permanently resident so every command buffer
+    /// the queue runs skips per-allocation residency tracking. Apple's
+    /// Metal driver otherwise re-validates a buffer's residency state at
+    /// each encode (the cost shows as host gap when there are tens of
+    /// thousands of small dispatches per prefill). One residency set
+    /// is shared across all weight buffers; subsequent calls add to it.
+    /// Requires macOS 15+ / iOS 18+. Opt-out via `FFAI_NO_RESIDENCY_SET=1`.
+    public func markWeightsResident(_ buffers: [MTLBuffer]) {
+        if ProcessInfo.processInfo.environment["FFAI_NO_RESIDENCY_SET"] != nil { return }
+        guard #available(macOS 15.0, iOS 18.0, *) else { return }
+        residencyLock.lock()
+        defer { residencyLock.unlock() }
+        if weightResidencySet == nil {
+            let descriptor = MTLResidencySetDescriptor()
+            descriptor.label = "FFAI weights"
+            descriptor.initialCapacity = max(buffers.count, 1024)
+            do {
+                let set = try mtlDevice.makeResidencySet(descriptor: descriptor)
+                commandQueue.addResidencySet(set)
+                weightResidencySet = set
+            } catch {
+                // Driver couldn't create the set — fall back silently.
+                weightResidencySet = nil
+                return
+            }
+        }
+        guard let set = weightResidencySet as? MTLResidencySet else { return }
+        for buf in buffers {
+            set.addAllocation(buf)
+        }
+        set.commit()
+        set.requestResidency()
+    }
 }

Sources/FFAI/Drafter.swift

@@ -0,0 +1,114 @@
+// Drafter — interface for speculative-decode candidate proposal.
+//
+// A drafter takes the token history so far and proposes one or more
+// candidate next tokens. The target model then verifies via
+// `forwardManyAllLogits([lastAccepted, ...candidates])` and either
+// accepts (commit candidates) or rejects (restore caches, commit the
+// model's actual choice from the verify logits).
+//
+// Implementations live alongside this protocol:
+//   * `NGramDrafter` — prompt-lookup n-gram. Zero ML cost.
+//   * `ANEMTPDrafter` — runs the trained MTP head on the ANE via the
+//     mlpackage built at /Users/tom/models/Qwen3.6-35B-A3B-mtp.mlpackage.
+//     (Implementation lands when the ANEDrafter Swift wrapper is wired.)
+//   * `GreedyDrafter` — stub that always returns nil. Useful for
+//     bench-comparing the spec-decode driver overhead vs raw decode.
+
+import Foundation
+
+/// A drafter proposes candidate next tokens for the target model to
+/// verify in a speculative-decode loop.
+public protocol Drafter: AnyObject {
+    /// Propose up to `gamma` candidate tokens. Implementations may
+    /// return fewer (or `nil` / empty) if they can't make a confident
+    /// proposal — the driver falls back to a plain decode step in
+    /// that case.
+    ///
+    /// `history` is the full sequence so far (prompt + generated).
+    /// `gamma` is the maximum candidates the driver is asking for.
+    func propose(history: [Int], gamma: Int) -> [Int]
+}
+
+/// Prompt-lookup n-gram drafter — zero ML cost; works well on
+/// repetitive contexts (code, structured chat).
+///
+/// Algorithm (matches the typical prompt-lookup decoding paper):
+///   1. Take the last `nMatch` tokens of `history` as the lookup key.
+///   2. Scan backwards through `history` looking for a previous
+///      occurrence of that key.
+///   3. If found, return the next `gamma` tokens AFTER that earlier
+///      occurrence — those are the candidate next tokens.
+///   4. If not found, fall back to shorter keys (`nMatch - 1`,
+///      `nMatch - 2`, ...) before giving up.
+///
+/// Common heuristic: start with `nMatch = 3` (trigrams) and fall back
+/// to bigrams + unigrams.
+public final class NGramDrafter: Drafter {
+    /// Largest match length to try first. Falls back to shorter lengths
+    /// if no longer match is found. Typical: 3 (trigram).
+    public let maxNMatch: Int
+    /// Smallest match length the drafter will try before giving up.
+    /// Default 1 (unigram lookup is usually too noisy; raising to 2
+    /// avoids spurious draft.)
+    public let minNMatch: Int
+
+    public init(maxNMatch: Int = 3, minNMatch: Int = 2) {
+        precondition(maxNMatch >= minNMatch && minNMatch >= 1,
+                     "NGramDrafter: maxNMatch (\(maxNMatch)) must be ≥ minNMatch (\(minNMatch)) ≥ 1")
+        self.maxNMatch = maxNMatch
+        self.minNMatch = minNMatch
+    }
+
+    public func propose(history: [Int], gamma: Int) -> [Int] {
+        precondition(gamma >= 0, "NGramDrafter.propose: gamma must be ≥ 0")
+        guard gamma > 0, !history.isEmpty else { return [] }
+
+        // Try longest match first, fall back to shorter. Returns the
+        // FIRST (most-recent) occurrence's continuation. High accept
+        // probability when the trigram repeats.
+        //
+        // Empirically (Qwen3.6-A3B + bench in SpecDecodeBenchTests):
+        // trigram acceptance ~83%, bigram ~50%, unigram-frequency ~35%.
+        // Spec-decode break-even acceptance at γ=2 with the current
+        // verify-cost ratio is ~70% — so unigram is a NET LOSS even
+        // though it increases proposal rate. Better to return [] and
+        // let the driver run a single decode step (52 ms) than waste
+        // a 77 ms batched verify on a low-confidence proposal.
+        for nMatch in stride(from: maxNMatch, through: minNMatch, by: -1) {
+            guard history.count >= nMatch else { continue }
+            let keyStart = history.count - nMatch
+            let key = Array(history[keyStart..<history.count])
+            // Scan backwards from just-before the key (so we don't
+            // match the key against itself).
+            var probe = keyStart - 1
+            while probe >= nMatch - 1 {
+                if matches(history, at: probe - (nMatch - 1), key: key) {
+                    let candidateStart = probe + 1
+                    let candidateEnd = Swift.min(candidateStart + gamma,
+                                                  history.count)
+                    if candidateEnd > candidateStart {
+                        return Array(history[candidateStart..<candidateEnd])
+                    }
+                }
+                probe -= 1
+            }
+        }
+        return []
+    }
+
+    @inline(__always)
+    private func matches(_ history: [Int], at start: Int, key: [Int]) -> Bool {
+        if start < 0 || start + key.count > history.count { return false }
+        for i in 0..<key.count {
+            if history[start + i] != key[i] { return false }
+        }
+        return true
+    }
+}
+
+/// Stub drafter that never proposes anything. Used to measure pure
+/// spec-decode driver overhead against the no-spec baseline.
+public final class NeverDrafter: Drafter {
+    public init() {}
+    public func propose(history: [Int], gamma: Int) -> [Int] { [] }
+}