diff --git a/Sources/FFAI/CacheSnapshot.swift b/Sources/FFAI/CacheSnapshot.swift new file mode 100644 index 00000000..6aec498a --- /dev/null +++ b/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) + } +} diff --git a/Sources/FFAI/ConvStateCache.swift b/Sources/FFAI/ConvStateCache.swift index 090bfe7c..9557d387 100644 --- a/Sources/FFAI/ConvStateCache.swift +++ b/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 { diff --git a/Sources/FFAI/Device.swift b/Sources/FFAI/Device.swift index ce1c57cd..7b04c0b9 100644 --- a/Sources/FFAI/Device.swift +++ b/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() + } } diff --git a/Sources/FFAI/Drafter.swift b/Sources/FFAI/Drafter.swift new file mode 100644 index 00000000..c9fb222a --- /dev/null +++ b/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..= 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.. Bool { + if start < 0 || start + key.count > history.count { return false } + for i in 0.. [Int] { [] } +} diff --git a/Sources/FFAI/GDNStateCache.swift b/Sources/FFAI/GDNStateCache.swift index 76ef62bf..ac1d6370 100644 --- a/Sources/FFAI/GDNStateCache.swift +++ b/Sources/FFAI/GDNStateCache.swift @@ -47,6 +47,12 @@ public final class GDNStateCache: LayerCacheProtocol, @unchecked Sendable { /// The two recurrent-state buffers, each shape `[Hv, Dv, Dk]`, fp32. /// `current` is the live state (kernel input); `next` receives the /// kernel's output. `swap()` exchanges them after each step. + /// + /// Both the legacy `Ops.gatedDeltaStep` path and the fused + /// `Ops.gatedDeltaPrepStep` path share these slots: the fused + /// kernel runs in fp32 (against bf16 model activations cast to + /// fp32 on the GPU first) so the state precision matches the + /// canonical legacy path exactly. public private(set) var current: Tensor public private(set) var next: Tensor @@ -74,11 +80,17 @@ public final class GDNStateCache: LayerCacheProtocol, @unchecked Sendable { self.next = Tensor.empty(shape: shape, dtype: .f32, device: device) self.current.zero() self.next.zero() + // GDN state buffers live for the lifetime of a generation. Pin + // them in the device residency set — every fused-GDN dispatch + // reads + writes them, so skipping per-encode residency + // tracking matters at 30 GDN layers × per-token cadence. + device.markWeightsResident([self.current.buffer, self.next.buffer]) } /// Exchange `current` and `next`. Call this after `Ops.gatedDeltaStep` - /// has written the updated state into `next`, so the next decode - /// step reads the fresh state via `current`. + /// or `Ops.gatedDeltaPrepStep` has written the updated state into + /// `next`, so the next decode step reads the fresh state via + /// `current`. Increments `length`. public func swap() { Swift.swap(¤t, &next) length += 1 @@ -92,12 +104,86 @@ public final class GDNStateCache: LayerCacheProtocol, @unchecked Sendable { length = 0 } + /// Snapshot the current state to a fresh `Tensor`. Used by + /// speculative-decode to remember the pre-speculative state so it + /// can be restored if drafted tokens are rejected. + /// + /// GDN is recurrent — once the kernel has folded a wrong token + /// into `current`, you cannot subtract it back. Snapshot before + /// each speculative forward, restore on full rejection. + /// + /// Snapshot cost: one `memcpy` of `[Hv, Dv, Dk]` fp32 = ~2 MiB on + /// Qwen3.6-A3B per layer × 30 layers = ~60 MiB. On Apple silicon + /// unified memory at ~400 GB/s this is ~150 µs per layer in + /// aggregate — negligible against a ~60 ms decode step. + /// + /// Returns the snapshot tensor; caller stores it until needed. + /// Cached snapshot tensor. Reused across snapshot calls so + /// spec-decode doesn't churn ~2 MB per layer per verify step. + private var cachedSnapshot: Tensor? + + public func snapshot(device: Device = .shared) -> Tensor { + let shape = [numValueHeads, valueHeadDim, keyHeadDim] + if cachedSnapshot == nil { + cachedSnapshot = Tensor.empty(shape: shape, dtype: .f32, device: device) + } + let snap = cachedSnapshot! + // Use a blit encoder to copy on the GPU side — keeps the host + // out of the loop. Both buffers are shared-storage, so the + // copy is just a unified-memory blit. + let cmd = device.makeCommandBuffer() + guard let blit = cmd.makeBlitCommandEncoder() else { + preconditionFailure("GDNStateCache.snapshot: makeBlitCommandEncoder failed") + } + let bytes = numValueHeads * valueHeadDim * keyHeadDim * DType.f32.byteSize + blit.copy(from: current.buffer, sourceOffset: current.offset, + to: snap.buffer, destinationOffset: snap.offset, + size: bytes) + blit.endEncoding() + cmd.commit() + cmd.waitUntilCompleted() + return snap + } + + /// Restore from a snapshot taken via `snapshot()`. Overwrites + /// `current` with the snapshot contents; leaves `next` alone (it'll + /// be overwritten on the next kernel dispatch). Does NOT decrement + /// `length` — caller is responsible for matching the position + /// counter to the restored state. + public func restore(from snapshot: Tensor, device: Device = .shared) { + let expected = numValueHeads * valueHeadDim * keyHeadDim + precondition(snapshot.elementCount == expected, + "GDNStateCache.restore: snapshot has \(snapshot.elementCount) elements, expected \(expected)") + precondition(snapshot.dtype == .f32, + "GDNStateCache.restore: snapshot must be f32 (matches state buffer dtype)") + let cmd = device.makeCommandBuffer() + guard let blit = cmd.makeBlitCommandEncoder() else { + preconditionFailure("GDNStateCache.restore: makeBlitCommandEncoder failed") + } + let bytes = expected * DType.f32.byteSize + blit.copy(from: snapshot.buffer, sourceOffset: snapshot.offset, + to: current.buffer, destinationOffset: current.offset, + size: bytes) + blit.endEncoding() + cmd.commit() + cmd.waitUntilCompleted() + } + /// Bytes physically allocated — both state buffers. Independent of /// how many tokens have been processed (the GDN compression win). public var bytesAllocated: Int { 2 * numValueHeads * valueHeadDim * keyHeadDim * DType.f32.byteSize } + /// Set the position counter directly without zeroing buffers. + /// Spec-decode restore path uses this after writing the snapshot + /// tensor into `current`. `reset()` + `swap()` would wipe the + /// just-restored state. + public func setLength(_ length: Int) { + precondition(length >= 0, "GDNStateCache.setLength: must be ≥ 0") + self.length = length + } + /// GDN state storage is constant-size; "in use" equals allocated /// once any step has executed, zero before then. public var bytesInUse: Int { diff --git a/Sources/FFAI/KVCache.swift b/Sources/FFAI/KVCache.swift index f433cc15..d7ca9b0b 100644 --- a/Sources/FFAI/KVCache.swift +++ b/Sources/FFAI/KVCache.swift @@ -174,6 +174,10 @@ public final class KVCache: KVCacheProtocol, @unchecked Sendable { self.kBuffer.zero() self.vBuffer.zero() self._evictionState = KVEvictionState(policy: eviction, bufferCapacity: maxSeq) + // KV buffers live for the lifetime of a generation — pin them in + // the device's residency set so per-dispatch residency tracking + // doesn't fire on the thousands of decode-step appends + reads. + device.markWeightsResident([self.kBuffer.buffer, self.vBuffer.buffer]) } /// CPU-side legacy append. Caller must have already sync'd the @@ -209,12 +213,13 @@ public final class KVCache: KVCacheProtocol, @unchecked Sendable { precondition(kFlat.dtype == dtype && vFlat.dtype == dtype, "KVCache: dtype mismatch") lengthLock.withLock { let pos = _evictionState.reserveNextSlot() - Ops.kvCacheUpdate(src: kFlat, into: kBuffer, - nKVHeads: nKVHeads, headDim: headDim, - maxSeq: maxSeq, position: pos, on: cmd) - Ops.kvCacheUpdate(src: vFlat, into: vBuffer, - nKVHeads: nKVHeads, headDim: headDim, - maxSeq: maxSeq, position: pos, on: cmd) + // ITER 10: shared encoder for K + V — saves 1 encoder + // begin/end pair per attn layer × 10 attn = ~170 µs/token. + Ops.kvCacheUpdateKV( + kSrc: kFlat, kCache: kBuffer, + vSrc: vFlat, vCache: vBuffer, + nKVHeads: nKVHeads, headDim: headDim, + maxSeq: maxSeq, position: pos, on: cmd) } } diff --git a/Sources/FFAI/Layers.swift b/Sources/FFAI/Layers.swift index 70a03928..30d1ee7d 100644 --- a/Sources/FFAI/Layers.swift +++ b/Sources/FFAI/Layers.swift @@ -41,6 +41,16 @@ public final class Linear: Module { } return y } + + /// T-batched. `x` is `[T, inDim]` flat row-major, returns + /// `[T, outDim]` flat. Dispatches `Ops.gemm` (one kernel, no per-row + /// host overhead). Bias path is unimplemented — the families that + /// currently call this (Qwen3-series) do not use linear biases. + public func callMany(_ x: Tensor, t: Int, on cmd: MTLCommandBuffer) -> Tensor { + precondition(bias == nil, + "Linear.callMany: bias broadcast over T rows not implemented; no Qwen3-series caller needs it today") + return Ops.gemm(weight: weight, input: x, nRows: t, on: cmd) + } } // ─── QuantizedLinear (mlx int4 format) ──────────────────────────────── @@ -79,11 +89,60 @@ public final class QuantizedLinear: Module { } public func callAsFunction(_ x: Tensor, on cmd: MTLCommandBuffer) -> Tensor { + // NOTE (ITER 7 revert): tried caching outScratch across decode + // steps to eliminate ~1300 Tensor.empty allocs per token, but + // bench showed slight regression (~1.5%). Hypothesis: Metal's + // hazard tracking added more fences across decode steps with + // shared scratch than the cost of fresh-buffer allocations + // (which MTLDevice likely pools internally). Fresh alloc per + // call is the verified faster path. Ops.dequantGemv( weight: weight, scales: scales, biases: biases, input: x, bits: bits, groupSize: groupSize, on: cmd ) } + + /// T-batched. `x` is `[T, inDim]` flat row-major, returns + /// `[T, outDim]` flat. Dispatches `Ops.dequantGemmDynamicM` — one + /// `mt_qmm_mma` kernel for T multiple of 32, or one padded + /// dispatch + host slice for ragged T. Matches the batched-prefill + /// path (BP2). + public func callMany(_ x: Tensor, t: Int, + on cmd: MTLCommandBuffer, + device: Device) -> Tensor { + let outDim = weight.shape[0] + let inDim = scales.shape[scales.shape.count - 1] * groupSize + // 4-bit goes through the fast mt_qmm_mma path. Other bit-widths + // (commonly 8-bit on smaller projections like Qwen3.5's + // shared_expert_gate at hidden→1) fall back to T sequential + // `dequantGemv` calls on the same `cmd`. Slower than a true + // batched kernel but bit-identical to the per-token path, and + // the projections that hit this branch are tiny so the per-token + // launch overhead is in the noise. + if bits == 4 { + let out = Tensor.empty(shape: [t, outDim], dtype: x.dtype, device: device) + Ops.dequantGemmDynamicM( + input: x, weight: weight, scales: scales, biases: biases, + t: t, nOut: outDim, kIn: inDim, groupSize: groupSize, + on: cmd, device: device, into: out) + return out + } + let dtBytes = x.dtype.byteSize + let out = Tensor.empty(shape: [t, outDim], dtype: x.dtype, device: device) + for r in 0.. Tensor + private let forwardMany: (Tensor, Int, MTLCommandBuffer, Device) -> Tensor public init(_ linear: Linear) { self.inner = linear self.forward = { linear($0, on: $1) } + self.forwardMany = { x, t, cmd, _ in linear.callMany(x, t: t, on: cmd) } } public init(_ linear: QuantizedLinear) { self.inner = linear self.forward = { linear($0, on: $1) } + self.forwardMany = { x, t, cmd, dev in + linear.callMany(x, t: t, on: cmd, device: dev) + } } public func parameters() -> [(String, Tensor)] { inner.parameters() } @@ -107,6 +171,15 @@ public final class AnyLinear: Module { public func callAsFunction(_ x: Tensor, on cmd: MTLCommandBuffer) -> Tensor { forward(x, cmd) } + + /// T-batched. `x` is `[T, inDim]` flat row-major, returns + /// `[T, outDim]` flat. Dispatches one `gemm` (dense) or one + /// `dequantGemmDynamicM` (quantized). Unblocks batched-prefill. + public func callMany(_ x: Tensor, t: Int, + on cmd: MTLCommandBuffer, + device: Device) -> Tensor { + forwardMany(x, t, cmd, device) + } } /// Derive the affine-quantization bit-width of one mlx-quantized tensor diff --git a/Sources/FFAI/Model.swift b/Sources/FFAI/Model.swift index 11c7f2d9..98764a7d 100644 --- a/Sources/FFAI/Model.swift +++ b/Sources/FFAI/Model.swift @@ -4,6 +4,7 @@ // surface. import Foundation +import Metal import Tokenizers public enum ModelError: Error, CustomStringConvertible { @@ -191,6 +192,15 @@ public enum ModelRegistry { availableCapabilities: Capability.textOnly.union([.visionIn]), vlModel: vlm) } + // Text-only Qwen3.5 / 3.6 checkpoints with a vestigial + // `vision_config` block. The architecture string is a Qwen3.5 + // MoE family name and we have a text-only loader — route to + // it instead of throwing the VL-not-integrated error. The + // vision tower is not used. + if let arch = config.architecture, Qwen35.architectures.contains(arch) { + return try loadQwen35(config: config, weights: weights, + options: options, device: device) + } // Other VL families — the FFAI vision foundation // (VisionEncoder, ImagePreprocessing, VLModel splice, // conv2d/patch_embed/rope_2d Ops) is in tree, but these @@ -767,6 +777,17 @@ public final class Model: @unchecked Sendable { let config = try ModelConfig.load(from: dir) Debug.log(.load, "config: arch=\(config.architecture ?? "?") model_type=\(config.modelType ?? "?") hidden=\(config.hiddenSize ?? 0) layers=\(config.numLayers ?? 0)") let bundle = try SafeTensorsBundle(directory: dir, device: device) + // Pin all weight buffers in GPU residency. Avoids + // per-command-buffer residency re-validation on every + // prefill / decode dispatch. macOS 15+ / iOS 18+; opt-out + // via FFAI_NO_RESIDENCY_SET=1. + var weightBuffers: [MTLBuffer] = [] + for file in bundle.files { + for entry in file.entries.values { + weightBuffers.append(entry.buffer) + } + } + device.markWeightsResident(weightBuffers) let loaded = try ModelRegistry.dispatchAndLoad( config: config, weights: bundle, options: options, device: device ) diff --git a/Sources/FFAI/Models/GraniteMoeHybrid.swift b/Sources/FFAI/Models/GraniteMoeHybrid.swift index 41a89153..f5d20e34 100644 --- a/Sources/FFAI/Models/GraniteMoeHybrid.swift +++ b/Sources/FFAI/Models/GraniteMoeHybrid.swift @@ -692,7 +692,7 @@ public final class GraniteMoeHybridDenseMLP: Module { func forward(_ xNorm: Tensor, cmd: MTLCommandBuffer) -> Tensor { let g = gateProj(xNorm, on: cmd) let u = upProj(xNorm, on: cmd) - let inner = Ops.mul(Ops.silu(g, on: cmd), u, on: cmd) + let inner = Ops.swiglu(gate: g, up: u, on: cmd) return downProj(inner, on: cmd) } } diff --git a/Sources/FFAI/Models/Jamba.swift b/Sources/FFAI/Models/Jamba.swift index 63186d93..d3586e0d 100644 --- a/Sources/FFAI/Models/Jamba.swift +++ b/Sources/FFAI/Models/Jamba.swift @@ -650,7 +650,7 @@ public final class JambaDenseMLP: Module { func forward(_ xNorm: Tensor, cmd: MTLCommandBuffer) -> Tensor { let g = gateProj(xNorm, on: cmd) let u = upProj(xNorm, on: cmd) - let inner = Ops.mul(Ops.silu(g, on: cmd), u, on: cmd) + let inner = Ops.swiglu(gate: g, up: u, on: cmd) return downProj(inner, on: cmd) } } diff --git a/Sources/FFAI/Models/LFM2.swift b/Sources/FFAI/Models/LFM2.swift index 5d1b3429..c9c1cd25 100644 --- a/Sources/FFAI/Models/LFM2.swift +++ b/Sources/FFAI/Models/LFM2.swift @@ -632,7 +632,7 @@ public final class LFM2MLP: Module { func forward(_ x: Tensor, on cmd: MTLCommandBuffer) -> Tensor { let gate = w1(x, on: cmd) let up = w3(x, on: cmd) - let inner = Ops.mul(Ops.silu(gate, on: cmd), up, on: cmd) + let inner = Ops.swiglu(gate: gate, up: up, on: cmd) return w2(inner, on: cmd) } } diff --git a/Sources/FFAI/Models/MoELayer.swift b/Sources/FFAI/Models/MoELayer.swift index c53c1891..83717c17 100644 --- a/Sources/FFAI/Models/MoELayer.swift +++ b/Sources/FFAI/Models/MoELayer.swift @@ -48,6 +48,7 @@ import Foundation import Metal +import MetalTileSwift // ─── Gating mode ───────────────────────────────────────────────────── @@ -124,22 +125,51 @@ public struct MoERouter: Sendable { "MoERouter.route: logits.count \(logits.count) ≠ nExperts \(nExperts)") switch gatingMode { + case .softmaxThenTopK where normTopKProb && expertBias == nil: + // Fast path: `softmaxThenTopK + normTopKProb=true` produces + // the same weights as `.topKThenSoftmax`. Proof: + // softmax(L)[i] = exp(L[i] - M) / Z, M = max(L), Z = Σ + // top-K of softmax probs by value desc = top-K of raw + // logits (softmax is monotonic on selection) + // weights_pre[k] = exp(L[idx_k] - M) / Z + // renorm: weights[k] = weights_pre[k] / Σ_k weights_pre + // = exp(L[idx_k]) / Σ_k exp(L[idx_k]) + // = softmax over just the K picked logits + // Skips the full softmax over `nExperts=256` (~500 ops + // including 256 exp + sum reduction + 256 div). Saves + // ~40 K host ops / token on Qwen3.6-A3B (40 MoE layers × + // ~1000 ops per route saved). Qwen3 / Qwen3.5 / Qwen3.6 MoE + // all hit this branch (`norm_topk_prob: true` in every + // shipped config). + // + // REQUIRES `expertBias == nil`: the proof's "top-K of raw + // logits == top-K of softmax probs" step breaks once an + // additive per-expert bias is applied AFTER softmax (the + // bias is not monotonic in the raw logit). LFM2-MoE supplies + // an `expert_bias`, so it falls through to the general + // `.softmaxThenTopK` path below, which selects on + // `softmax(logits) + bias`. + let idx = Self.topKIndices(logits, k: topK) + let weights = Self.softmax(idx.map { logits[$0] }) + return Routing(indices: idx, weights: weights) + case .softmaxThenTopK: - // 1. softmax over ALL experts (numerically stable). + // `normTopKProb=false` case — semantically distinct + // (unnormalised softmax probs as weights). Must walk the + // full softmax over all nExperts. let probs = Self.softmax(logits) - // 2. optional per-expert additive bias (LFM2-MoE - // `expert_bias`). With no bias `gated == probs`, so the - // other softmaxThenTopK families are unaffected. + // Optional per-expert additive bias (LFM2-MoE `expert_bias`). + // With no bias `gated == probs`, so other softmaxThenTopK + // families are byte-identical. let gated: [Float] if let bias = expertBias { gated = zip(probs, bias).map { $0 + $1 } } else { gated = probs } - // 3. top-K of the (biased) gate values. let idx = Self.topKIndices(gated, k: topK) var weights = idx.map { gated[$0] } - // 4. optional re-normalisation of the K picked weights. + // Optional re-normalisation of the K picked weights. if normTopKProb { let sum = weights.reduce(0, +) if sum > 0 { weights = weights.map { $0 / sum } } @@ -147,10 +177,10 @@ public struct MoERouter: Sendable { return Routing(indices: idx, weights: weights) case .topKThenSoftmax: - // 1. top-K of the raw logits. + // Top-K of raw logits, softmax over just the K picked — + // always normalised by construction so `normTopKProb` + // doesn't apply. let idx = Self.topKIndices(logits, k: topK) - // 2. softmax over just the K picked logits — always - // normalised, so `normTopKProb` does not apply. let weights = Self.softmax(idx.map { logits[$0] }) return Routing(indices: idx, weights: weights) } @@ -173,13 +203,76 @@ public struct MoERouter: Sendable { /// Indices of the K largest values, ordered largest-first. Ties /// resolved by smaller index (matches argpartition semantics). static func topKIndices(_ x: [Float], k: Int) -> [Int] { - // Sort indices by (value desc, index asc). nExperts ≤ 128 so a - // full sort is cheaper than maintaining a heap. - let order = (0.. x[b] } - return a < b + // Partial-sort: maintain a k-sized min-heap of (value, index) + // pairs, ordered by (value asc, index desc) so the smallest + // value with the largest index is at the top (i.e. easiest to + // evict). At the end, drain the heap and reverse — gives the + // top-K ordered by (value desc, index asc), matching the prior + // full-sort semantics. + // + // Complexity: O(n log k) vs the prior `Array.sorted` O(n log n). + // At Qwen3.6-A3B (nExperts=256, k=8) this is ~280 ops vs ~2048 + // ops per call → ~7× host CPU per route × 40 MoE layers per + // decode token ≈ 5-7 ms / token saved. Stale comment claimed + // "nExperts ≤ 128 so a full sort is cheaper" — at 256 it + // already wasn't, and partial-sort is the right structure + // regardless. + precondition(k > 0 && k <= x.count, "topKIndices: k must be in 1...n") + var heap: [(value: Float, index: Int)] = [] + heap.reserveCapacity(k) + // Min-heap ordering: smallest value with largest index at root. + // Returns true when `a` is "smaller" (closer to root / first to + // evict). For tie on value, the entry with the LARGER index is + // smaller (we want smaller indices to win ties on the final + // sort, so the larger-index entry is the easier evict). + @inline(__always) func less(_ a: (Float, Int), _ b: (Float, Int)) -> Bool { + if a.0 != b.0 { return a.0 < b.0 } + return a.1 > b.1 + } + @inline(__always) func siftUp(_ start: Int) { + var i = start + while i > 0 { + let parent = (i - 1) >> 1 + if less(heap[i], heap[parent]) { + heap.swapAt(i, parent) + i = parent + } else { return } + } + } + @inline(__always) func siftDown(_ start: Int) { + var i = start + let n = heap.count + while true { + let l = 2 * i + 1 + let r = 2 * i + 2 + var smallest = i + if l < n && less(heap[l], heap[smallest]) { smallest = l } + if r < n && less(heap[r], heap[smallest]) { smallest = r } + if smallest == i { return } + heap.swapAt(i, smallest) + i = smallest + } } - return Array(order.prefix(k)) + for i in 0.. b.value } + return a.index < b.index + } + return heap.map { $0.index } } } @@ -212,18 +305,135 @@ public final class MoELayer: Module, DecoderLayer { public let router: MoERouter public let hidden: Int + /// Stacked-expert weight handles for the batched gather BGEMM fast + /// path. When set, `decode` dispatches `mt_moe_gather_qmm_mma_int4_bm16` + /// instead of running `topK` sequential per-expert SwiGLU triplets — + /// 3 kernel launches per token instead of `3 * topK` (24 → 3 at + /// `topK=8`). Populated by the host model's MoE builder when (a) the + /// checkpoint is mlx int4-quantized, (b) the stacked weight tensors + /// are intact (not sliced away by the per-expert wrapper path), and + /// (c) the moeIntermediate / hidden shapes satisfy the bm16 kernel's + /// `N % 32 == 0` and `K % 32 == 0` tile contract. Left nil for + /// everyone else — the legacy serial-expert loop still runs. + public struct StackedInt4Experts: Sendable { + /// `[numExperts, moeIntermediate, hidden/8]` u32 packed. + public let gateWeight: Tensor + /// `[numExperts, moeIntermediate, hidden/groupSize]` in `dtype`. + public let gateScales: Tensor + public let gateBiases: Tensor + /// `[numExperts, moeIntermediate, hidden/8]` u32 packed. + public let upWeight: Tensor + public let upScales: Tensor + public let upBiases: Tensor + /// `[numExperts, hidden, moeIntermediate/8]` u32 packed. + public let downWeight: Tensor + public let downScales: Tensor + public let downBiases: Tensor + public let numExperts: Int + public let moeIntermediate: Int + public let hidden: Int + public let groupSize: Int + /// Activation dtype the scales / biases / activations all use. + public let dtype: DType + + public init(gateWeight: Tensor, gateScales: Tensor, gateBiases: Tensor, + upWeight: Tensor, upScales: Tensor, upBiases: Tensor, + downWeight: Tensor, downScales: Tensor, downBiases: Tensor, + numExperts: Int, moeIntermediate: Int, hidden: Int, + groupSize: Int, dtype: DType) { + self.gateWeight = gateWeight + self.gateScales = gateScales + self.gateBiases = gateBiases + self.upWeight = upWeight + self.upScales = upScales + self.upBiases = upBiases + self.downWeight = downWeight + self.downScales = downScales + self.downBiases = downBiases + self.numExperts = numExperts + self.moeIntermediate = moeIntermediate + self.hidden = hidden + self.groupSize = groupSize + self.dtype = dtype + } + } + public let stackedInt4Experts: StackedInt4Experts? + + /// Env-flag cache — read once at init so the decode loop doesn't + /// pay 3 × `ProcessInfo.processInfo.environment[...]` dictionary + /// lookups per MoE layer per token. At Qwen3.6-A3B that's 120 + /// lookups / token saved (40 layers × 3 envs). Each lookup is + /// ~100 ns on Foundation, so the total is sub-noise (~12 µs / token) + /// — shipped for cleanliness more than the µbench delta. + public let enableBGEMM: Bool + public let useBm8Env: Bool + public let useM1Env: Bool + + /// Lazy cache of the dequantized gate weight in `[outDim=nExperts, + /// inDim=hidden]` layout — what `Ops.gemm` expects for its `weight` + /// argument (`ffai_gemm` computes `out = input @ weight^T`). For + /// 8-bit gates, `QuantizedLinear.callMany` falls into a T-loop of + /// per-row `dequantGemv` (T·40 = 20480 launches at T=512 prefill); + /// the dense path collapses each layer's gate to one `Ops.gemm` + /// dispatch regardless of T (no T % 64 alignment constraint). + /// Skipped when: + /// • `gate.inner` is not `QuantizedLinear`, or + /// • bits != 8, or + /// • env `FFAI_MOE_NO_DEQUANT_GATE=1`. + /// Materialised lazily on first `gateLogitsMany` to avoid eager work + /// in layers that never see a forward pass. + private var dequantizedGateCache: Tensor? + /// Sibling cache in `[hidden, nExperts]` (K, N) layout for the + /// `mt_steel_gemm_64x64x16_2x2_*` fast path. Engaged when T and + /// nExperts both divide 64 — covers the production prefill shape + /// (T=64, 128, 512; nExperts=128). + private var dequantizedGateCacheKN: Tensor? + private var dequantizedGateAttempted = false + + /// Cached scratch buffers for the per-expert weighted-sum loop at + /// decode T=1. Reused across decode calls — avoids 1 host alloc + + /// 1 host memset per layer per token (~10 µs × 40 layers = 400 µs + /// per Qwen3.6-A3B decode token). + /// * `accumulatorScratch`: [hidden] tensor; zeroed at the start of + /// each decode call, then accumulated into via `Ops.scalarFMA`. + /// * `topKScalarsBuf`: small [topK × dtype.byteSize] MTLBuffer + /// holding the 8 routing weights packed in the model dtype. + private var accumulatorScratch: Tensor? + private var topKScalarsBuf: MTLBuffer? + /// ITER 32 (post-OOM): per-expert g/u scratches cached at init. + /// Replaces the per-call `Tensor.empty([moeIntermediate])` in the + /// swiGLU helper that caused 640 fresh tensor allocations per + /// Qwen3.6-A3B decode token (8 experts × 40 layers × 2 = 640). + /// Each [moeIntermediate] bf16 = 2 KiB; 16 KiB per layer total. + /// Safe across decode tokens because the caller commits + waits + /// the layer's cmd before the next forward starts. + /// Indexed by slot (0..topK-1). + private var expertGScratches: [Tensor] = [] + private var expertUScratches: [Tensor] = [] + /// ITER 36: cache swiglu inner output per expert slot. + private var expertInnerScratches: [Tensor] = [] + /// ITER 38: cached per-expert down output for the batched 8-expert + /// down qmm path. Indexed by slot. + private var expertOutScratches: [Tensor] = [] + /// ITER 42: cached output of dense gate gemv (ITER 15). + private var gateLogitsScratch: Tensor? + /// - gate: hidden → nExperts router projection. /// - gateProj/upProj/downProj: `nExperts`-long arrays of per-expert /// SwiGLU projections, index-aligned with the expert id. /// - sharedGate/Up/DownProj: optional shared-expert SwiGLU; pass all /// three or none. /// - router: the top-K + gating-math configuration. + /// - stackedInt4Experts: optional batched-BGEMM fast path. Per-expert + /// arrays above remain the source of truth for `parameters()` + /// (checkpoint binding) and the fallback decode path. public init(gate: AnyLinear, gateProj: [AnyLinear], upProj: [AnyLinear], downProj: [AnyLinear], sharedGateProj: AnyLinear? = nil, sharedUpProj: AnyLinear? = nil, sharedDownProj: AnyLinear? = nil, - router: MoERouter, hidden: Int) { + router: MoERouter, hidden: Int, + stackedInt4Experts: StackedInt4Experts? = nil) { precondition(gateProj.count == router.nExperts, "MoELayer: gateProj has \(gateProj.count) experts, router expects \(router.nExperts)") precondition(upProj.count == router.nExperts, @@ -243,6 +453,19 @@ public final class MoELayer: Module, DecoderLayer { self.sharedDownProj = sharedDownProj self.router = router self.hidden = hidden + if let s = stackedInt4Experts { + precondition(s.numExperts == router.nExperts, + "MoELayer: stackedInt4Experts.numExperts \(s.numExperts) ≠ router.nExperts \(router.nExperts)") + precondition(s.hidden == hidden, + "MoELayer: stackedInt4Experts.hidden \(s.hidden) ≠ MoELayer.hidden \(hidden)") + precondition(s.moeIntermediate % 32 == 0 && hidden % 32 == 0, + "MoELayer: stackedInt4Experts shape (moeIntermediate=\(s.moeIntermediate), hidden=\(hidden)) violates bm16 N%32 / K%32 tile contract") + } + self.stackedInt4Experts = stackedInt4Experts + let env = ProcessInfo.processInfo.environment + self.enableBGEMM = env["FFAI_MOE_BGEMM"] != nil + self.useBm8Env = env["FFAI_MOE_BGEMM_BM8"] != nil + self.useM1Env = env["FFAI_MOE_M1"] != nil } public func parameters() -> [(String, Tensor)] { @@ -282,8 +505,33 @@ public final class MoELayer: Module, DecoderLayer { "MoELayer.decode: input has \(h.elementCount) elements, expected hidden \(hidden)") // ── 1. Gate gemv on the caller's command buffer ────────────── - // Queued onto `cmd` so it runs after whatever produced `h`. - let logitsTensor = gate(h, on: cmd) + // ITER 15: if we've already dequantized the gate weight (via a + // prior prefill call to gateLogitsMany), use the dense fp16 + // gemv path — bypasses the slow 8-bit dequant-gemv at T=1. + // For pure decode without prefill, this opportunistically + // dequantizes once per MoE layer (40 layers × hidden=2048 × + // nExperts=128 × 2 bytes = ~21 MB total — done once per + // layer's first decode call). + let logitsTensor: Tensor + if !dequantizedGateAttempted { + dequantizedGateAttempted = true + if ProcessInfo.processInfo.environment["FFAI_MOE_NO_DEQUANT_GATE"] == nil, + let q = gate.inner as? QuantizedLinear, + q.bits == 8 { + dequantizedGateCache = Self.dequantizeQuantizedLinear(q, device: device) + } + } + if let dense = dequantizedGateCache { + if gateLogitsScratch == nil + || gateLogitsScratch!.elementCount != router.nExperts + || gateLogitsScratch!.dtype != h.dtype { + gateLogitsScratch = Tensor.empty(shape: [router.nExperts], dtype: h.dtype) + } + logitsTensor = Ops.gemv(weight: dense, input: h, on: cmd, + into: gateLogitsScratch) + } else { + logitsTensor = gate(h, on: cmd) + } // ── 2. Commit + wait so the router can read the logits ─────── // The decode path batches a token onto one command buffer; the @@ -301,44 +549,889 @@ public final class MoELayer: Module, DecoderLayer { // contribute zero — skipping them is the same result, cheaper). let work = device.makeCommandBuffer() var accumulator: Tensor? - for (slot, expertId) in routing.indices.enumerated() { - // Broadcast the CPU combine weight into a [hidden] constant - // tensor so the element-wise `Ops.mul` can scale the expert - // output — avoids a dedicated scalar-multiply kernel. - let weightTensor = Tensor.filled(routing.weights[slot], - shape: [hidden], dtype: h.dtype, - device: device) - let expertOut = swiGLU(h, - gateProj: gateProj[expertId], - upProj: upProj[expertId], - downProj: downProj[expertId], - on: work) - let scaled = Ops.mul(expertOut, weightTensor, on: work) - accumulator = accumulator.map { Ops.add($0, scaled, on: work) } ?? scaled + // Batched gather-BGEMM fast path is opt-in via FFAI_MOE_BGEMM=1. + // At T=1 decode with topK=8 / m_total=8 it MEASURABLY REGRESSES + // vs the sequential per-expert matvec path on M5 Max — the + // bm16 tile pads to 16 rows but only 8 commit, so the kernel + // pays full weight-reload bandwidth for 50% useful work. Bench + // shows ~2.1× slowdown at decode T=1 on Qwen3.6-A3B + // (Iter 9: 20.34 tps with BGEMM → 43.20 tps without). + // + // The env flag stays opt-in for the prefill path (decodeMany) + // where mTotal scales with T and BGEMM wins, but at T=1 we + // hard-disable to avoid the regression even if the env is set. + // Override: `FFAI_MOE_BGEMM_FORCE_T1=1` re-enables for the rare + // case someone wants to bench the regression directly. + let forceBGEMMAtT1 = ProcessInfo.processInfo.environment["FFAI_MOE_BGEMM_FORCE_T1"] != nil + let useBGEMMAtT1 = enableBGEMM && forceBGEMMAtT1 + if let stacked = stackedInt4Experts, stacked.dtype == h.dtype, useBGEMMAtT1 { + // Fast path: one batched gather BGEMM per projection. The + // kernel expects rows of activations sorted by expert id; we + // sort the topK indices ascending and replicate `h` into the + // gathered row order. The per-row expert assignment goes in + // an int32 indices buffer. + accumulator = batchedSwiGLU(h, stacked: stacked, routing: routing, + on: work, device: device) + } else { + // Pack the topK routing weights into a single [topK] scalar + // buffer once, then dispatch one `scalar_fma` per expert + // instead of the old `Tensor.filled([hidden]) + Ops.mul + + // Ops.add` chain. Saves a host alloc + 2 dispatches per + // expert; at Qwen3.6-A3B topK=8 × 40 layers that's 320 + // allocations + 640 dispatches per decode token. + // + // First iteration: bootstrap the accumulator by allocating + // [hidden] zeros so scalar_fma can read it as base. The + // zero buffer itself is cheap (1 dispatch on a small tensor + // OR one Tensor.empty + .zero); we hold it until the loop + // is done. Subsequent iterations alias `accumulator` as + // both base and out (safe — the kernel reads at idx then + // writes at idx). + // Reuse cached scratch slots — see field comment on + // `accumulatorScratch` + `topKScalarsBuf`. + let scalarBufBytes = routing.weights.count * h.dtype.byteSize + if topKScalarsBuf == nil || topKScalarsBuf!.length < scalarBufBytes { + topKScalarsBuf = device.makeBuffer(length: scalarBufBytes) + } + Self.writeTopKScalars(routing.weights, dtype: h.dtype, + into: topKScalarsBuf!) + let scalarBuf = topKScalarsBuf! + + // Lazy-init + zero the accumulator scratch. + if accumulatorScratch == nil + || accumulatorScratch!.dtype != h.dtype + || accumulatorScratch!.elementCount != hidden { + accumulatorScratch = Tensor.empty(shape: [hidden], dtype: h.dtype, + device: device) + } + accumulatorScratch!.zero() + let acc = accumulatorScratch! + + // ITER 38: phase-2 batched down qmm. When all 8 expert + // downProjs are int4 QuantizedLinear with same groupSize, + // dispatch all 8 down qmms in ONE encoder. Uses cached + // inner + out scratches (no allocation per call). + var canBatchDown = true + var downWeights: [Tensor] = [] + var downScales: [Tensor] = [] + var downBiases: [Tensor] = [] + var downGroupSize: Int = 0 + for expertId in routing.indices { + if let q = downProj[expertId].inner as? QuantizedLinear, q.bits == 4 { + if downGroupSize == 0 { downGroupSize = q.groupSize } + if q.groupSize != downGroupSize { canBatchDown = false; break } + downWeights.append(q.weight) + downScales.append(q.scales) + downBiases.append(q.biases) + } else { + canBatchDown = false + break + } + } + + if canBatchDown { + // Phase 1a: per-expert gate+up batched (ITER 25), each + // expert is one shared-encoder dispatch of 2 qmms. + for (slot, expertId) in routing.indices.enumerated() { + let qg = gateProj[expertId].inner as! QuantizedLinear + let qu = upProj[expertId].inner as! QuantizedLinear + while expertGScratches.count <= slot { + let outDim = qg.weight.shape[0] + expertGScratches.append(Tensor.empty(shape: [outDim], dtype: h.dtype)) + expertUScratches.append(Tensor.empty(shape: [outDim], dtype: h.dtype)) + expertInnerScratches.append(Tensor.empty(shape: [outDim], dtype: h.dtype)) + } + while expertOutScratches.count <= slot { + expertOutScratches.append(Tensor.empty(shape: [hidden], dtype: h.dtype)) + } + Ops.dequantGemvInt4Two( + input: h, + w0: qg.weight, s0: qg.scales, b0: qg.biases, + out0: expertGScratches[slot], + w1: qu.weight, s1: qu.scales, b1: qu.biases, + out1: expertUScratches[slot], + groupSize: qg.groupSize, on: work) + } + // Phase 1b (ITER 39): batched 8-expert swiglu in ONE + // encoder using cached scratches. Saves 7 begin/end + // pairs per layer × 40 = 280/token. + let gs = Array(expertGScratches.prefix(routing.indices.count)) + let us = Array(expertUScratches.prefix(routing.indices.count)) + let innersOut = Array(expertInnerScratches.prefix(routing.indices.count)) + Ops.swigluMany(gates: gs, ups: us, outs: innersOut, on: work) + // Phase 2: batched 8-expert down qmm in ONE encoder. + let inners = Array(expertInnerScratches.prefix(routing.indices.count)) + let outs = Array(expertOutScratches.prefix(routing.indices.count)) + Ops.dequantGemvInt4Many( + weights: downWeights, scales: downScales, biases: downBiases, + inputs: inners, outputs: outs, + groupSize: downGroupSize, on: work) + // Phase 3: scalarFMA accumulate. + for (slot, expertOut) in outs.enumerated() { + let scalarT = Tensor(buffer: scalarBuf, + offset: slot * h.dtype.byteSize, + shape: [1], dtype: h.dtype) + Ops.scalarFMA(scalar: scalarT, value: expertOut, base: acc, + into: acc, on: work) + } + } else { + // Legacy: per-expert swiGLU chain. + for (slot, expertId) in routing.indices.enumerated() { + let expertOut = swiGLU(h, slot: slot, + gateProj: gateProj[expertId], + upProj: upProj[expertId], + downProj: downProj[expertId], + on: work) + let scalarT = Tensor(buffer: scalarBuf, + offset: slot * h.dtype.byteSize, + shape: [1], dtype: h.dtype) + Ops.scalarFMA(scalar: scalarT, value: expertOut, base: acc, + into: acc, on: work) + } + } + accumulator = acc } // ── 5. Optional always-on shared expert ────────────────────── if let sg = sharedGateProj, let su = sharedUpProj, let sd = sharedDownProj { - let sharedOut = swiGLU(h, gateProj: sg, upProj: su, downProj: sd, on: work) + // Shared expert uses slot = topK to avoid collision with + // per-expert scratches (e.g., slot 8 when topK = 8). + let sharedOut = swiGLU(h, slot: router.topK, + gateProj: sg, upProj: su, downProj: sd, on: work) accumulator = accumulator.map { Ops.add($0, sharedOut, on: work) } ?? sharedOut } // topK ≥ 1 so `accumulator` is always non-nil here. let result = accumulator! + // Commit without wait: `result` is on the in-flight `work` + // buffer; the caller's residual-add cmd will hazard-track the + // read against this write. Saves ~0.5-1 ms host-stall per + // MoE layer per token (Qwen3.6-A3B = 40 layers). work.commit() - work.waitUntilCompleted() return result } + /// T-batched MoE forward. `hFlat` is `[T, hidden]` flat; returns + /// `[T, hidden]` flat. The mTotal = T·topK rows fan into a single + /// `Ops.moeGatherDequantGemmInt4` (or bm8 variant) per projection + /// instead of T·topK sequential per-expert SwiGLU triplets. At + /// Qwen3.6-A3B T=32, topK=8 → mTotal=256 fills the BM=16 tiles + /// 16-deep — the regime where cooperative-tensor weight sharing + /// PAYS (the same kernels regress at T=1 because every row is a + /// unique expert with no shared weight). + /// + /// Requires `stackedInt4Experts` (the BGEMM weights layout). Falls + /// back to per-token `decode` loop otherwise. + /// + /// Architecture: + /// 1. Gate gemv batched (`gate.callMany`) → `[T, nExperts]`. + /// 2. `cmd.commit() + wait` — router needs logits on CPU. + /// 3. Per-token routing on host (T calls, each O(nExperts) + a + /// partial sort). + /// 4. Build plan over T·topK rows: `sourceToken[m]`, + /// `expertId[m]`, `weight[m]`. Sort by `expertId` ascending so + /// consecutive rows share an expert (the kernel walks rows in + /// tile order; sorted = weight tile reuse across rows in a + /// tile). + /// 5. `Ops.gather(h, gatherIdx)` → `xGathered [mTotal, hidden]` + /// in one dispatch (vs T·topK host memcpys). + /// 6. `moeGatherDequantGemmInt4` × 3 (gate / up / down) on the + /// gathered batch — one dispatch each. + /// 7. `Ops.swiglu` element-wise across `[mTotal, moeIntermediate]`. + /// 8. Weighted scatter-sum: for each `m`, scale `downOut[m]` by + /// `weights[m]` and accumulate into `outFlat[sourceToken[m]]`. + /// Uses the same `Tensor.filled([hidden])` broadcast pattern as + /// the single-token `batchedSwiGLU`, mTotal times. Scatter-sum + /// is the dominant residual cost at large `mTotal`; a future + /// `mt_moe_scatter_scale_add` kernel collapses these mTotal + /// dispatches into one. + /// 9. Shared expert (if any): one per-row SwiGLU call across T + /// tokens (also currently a T-loop because the single-token + /// `swiGLU` path is the standing API). + /// + /// All work after the gate-readback runs on a fresh `work` cmd that + /// commits once at the end. The caller's `cmd` is consumed by the + /// gate gemv and the commit + wait. + public func decodeMany(_ hFlat: Tensor, t: Int, + cmd: MTLCommandBuffer, device: Device) -> Tensor { + precondition(hFlat.elementCount == t * hidden, + "MoELayer.decodeMany: hFlat size \(hFlat.elementCount) ≠ T·hidden = \(t * hidden)") + precondition(t > 0, "MoELayer.decodeMany: T must be positive") + + let dt = hFlat.dtype + + // ── 1. Gate gemv batched on caller's cmd ───────────────────────── + let hRows = hFlat.reshaped(to: [t, hidden]) + let gateLogitsAll = gateLogitsMany(hRows, t: t, on: cmd, device: device) + // gateLogitsAll shape: [T, nExperts] + + // ── 2. Commit + wait so the router can read logits ─────────────── + cmd.commit() + cmd.waitUntilCompleted() + + // ── 3. Per-token routing on host ───────────────────────────────── + // `router.route` is a pure function over a per-row logit slice — + // independent across `r`, so the T calls fan out across CPU cores + // via `DispatchQueue.concurrentPerform`. At Qwen3.6-A3B T=512 × + // 40 layers = 20 480 route() calls per prefill; serial they add + // up. Pre-allocate the result array sized to `t` so each iteration + // writes its own slot (race-free). + let nExperts = router.nExperts + let topK = router.topK + let logitsHost = gateLogitsAll.toFloatArray() // [T·nExperts] + var routings = [MoERouter.Routing](repeating: MoERouter.Routing(indices: [], weights: []), + count: t) + routings.withUnsafeMutableBufferPointer { buf in + DispatchQueue.concurrentPerform(iterations: t) { r in + let start = r * nExperts + let rowLogits = Array(logitsHost[start..<(start + nExperts)]) + buf[r] = router.route(logits: rowLogits) + } + } + + // ── 4. Build sorted plan over T·topK rows ──────────────────────── + let mTotal = t * topK + // Tuple: (sortKey expertId, sourceToken, originalSlot, weight) + var planTuples: [(Int, Int, Int, Float)] = [] + planTuples.reserveCapacity(mTotal) + for r in 0..