diff --git a/Sources/FFAI/Device.swift b/Sources/FFAI/Device.swift index c63d9899..44b2ac59 100644 --- a/Sources/FFAI/Device.swift +++ b/Sources/FFAI/Device.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -24,6 +24,16 @@ public final class Device: @unchecked Sendable { public let mtlDevice: MTLDevice public let commandQueue: MTLCommandQueue + /// Lazy MTLResidencySet that pins the model's weight buffers so + /// every command buffer skips per-allocation residency tracking. + /// Populated after `Model.load` finishes. Typed as `Any?` so the + /// deployment target stays below macOS 15; the cast back to + /// `MTLResidencySet` lives inside the `@available` block in + /// `markWeightsResident`. Initialised under `residencyLock` to + /// single-flight the descriptor build. + 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. @@ -51,4 +61,40 @@ public final class Device: @unchecked Sendable { } return cb } + + /// Add `buffers` to a persistent MTLResidencySet attached to the + /// command queue. Without this, Apple's Metal driver re-validates + /// per-allocation residency on every command-buffer encode — at + /// model sizes with tens of thousands of dispatches per prefill, + /// the per-dispatch overhead dominates wall time. One residency + /// set is shared across all weight buffers; repeated calls add + /// to it. Requires macOS 15+ / iOS 18+; older OSes silently + /// no-op. Set `FFAI_NO_RESIDENCY_SET=1` to disable for A/B. + 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 refused to create the set; fall back to default + // residency tracking. Not fatal — just slower. + 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/Generation/Drafter.swift b/Sources/FFAI/Generation/Drafter.swift new file mode 100644 index 00000000..5443754c --- /dev/null +++ b/Sources/FFAI/Generation/Drafter.swift @@ -0,0 +1,471 @@ +// Copyright 2026 Tom Turney (@TheTom) +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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 a batched +// forward over `[lastAccepted, ...candidates]` and either accepts +// (commit candidates) or rejects (restore caches, commit the model's +// actual choice from the verify logits). +// +// Implementations: +// * `NGramDrafter` — prompt-lookup n-gram. Zero ML cost; works well +// on repetitive contexts (code, structured chat). +// * `NeverDrafter` — stub that always returns nothing. Used to +// measure pure spec-decode driver overhead vs raw decode. +// * `NGramTreeDrafter` — real branching n-gram tree drafter for +// tree-verify spec-decode (top-K continuations per depth). + +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 an empty array) 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: +/// 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. +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 2 (unigram lookup is too noisy in practice and pushes + /// average acceptance below the γ=2 break-even of ~70%). + 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 [] } + + 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] { [] } +} + +// ─── Tree drafter (tree-verify spec decode) ────────────────────────── +// +// Linear γ=2 spec decode caps at 1.83 expected accepted tokens per +// verify cycle (83% acceptance × 2 candidates + 1 verified). Tree +// drafters expand multiple continuations per draft step, verifying +// ALL of them in one forward pass with a tree-causal attention mask. +// At γ=8 tree (e.g., 2-branch at depth 3) with 50–65% per-branch +// acceptance, expected accepted tokens jumps to 3–4 → 1.7–2× decode +// over linear. + +/// One node in a draft tree. The root represents the first proposed +/// token following `history`; each child represents a possible +/// continuation after its parent. A linear chain (one child per node) +/// degenerates to the existing `[Int]` candidate sequence. +public struct DraftTreeNode: Sendable { + /// The token this node proposes. + public let token: Int + /// Continuations after this token. Empty at leaves. + public let children: [DraftTreeNode] + + public init(token: Int, children: [DraftTreeNode] = []) { + self.token = token + self.children = children + } + + /// Total node count (root + all descendants). + public var size: Int { + return 1 + children.reduce(0) { $0 + $1.size } + } + + /// Maximum depth (root alone = 1). + public var depth: Int { + return 1 + (children.map(\.depth).max() ?? 0) + } + + /// Depth-first flatten of the tree. Returns: + /// - `tokens[i]`: the token at flat position `i` (root at 0). + /// - `parentIndex[i]`: the flat-index of node `i`'s parent, or + /// `-1` for the root (`i == 0`). + /// - `pathFromRoot[i]`: indices of node `i`'s ancestors INCLUDING + /// `i` itself, root-first (length = node `i`'s depth in the + /// tree). Used by the verify driver to walk the accepted prefix. + public func flatten() -> ( + tokens: [Int], parentIndex: [Int], pathFromRoot: [[Int]] + ) { + var tokens: [Int] = [] + var parent: [Int] = [] + var paths: [[Int]] = [] + var indexStack: [Int] = [] // DFS path-to-this-node, by flat-index + func recurse(_ node: DraftTreeNode, parentIdx: Int) { + let myIdx = tokens.count + tokens.append(node.token) + parent.append(parentIdx) + indexStack.append(myIdx) + paths.append(indexStack) + for child in node.children { + recurse(child, parentIdx: myIdx) + } + indexStack.removeLast() + } + recurse(self, parentIdx: -1) + return (tokens, parent, paths) + } + + /// Result of walking a draft tree against the target's argmax + /// predictions. + public struct VerifyResult: Equatable, Sendable { + /// The accepted-path tokens — root + each child whose token + /// matched the target's argmax at the prior depth, inclusive. + /// Empty if the root didn't match the target's first prediction. + public let acceptedTokens: [Int] + /// The "bonus" token from the target's argmax at the deepest + /// accepted flat-position (or at the pre-tree position if + /// even the root failed). This is the standard spec-decode + /// guaranteed-correct-token-for-free. + public let bonusToken: Int + } + + /// Walk the tree against the target's per-position argmax oracle; + /// return the longest accepted path + the bonus token. + /// + /// Protocol: + /// 1. `oracleAtHistoryEnd` is the target's argmax of logits at + /// the position JUST BEFORE the tree (the last token in + /// history). It's the target's preferred root token. If it + /// doesn't equal `self.token`, the root is rejected — return + /// no accepted tokens + `oracleAtHistoryEnd` as the bonus. + /// 2. Otherwise the root is accepted; descend by repeatedly + /// looking up `oracle(currentFlatIndex)` and accepting + /// whichever child has that token. Stop when no child + /// matches; return accepted path + `oracle(lastAcceptedIdx)` + /// as the bonus. + /// + /// Pure function. No model / cache dependencies. The driver wires + /// this up with `oracle = { i in target_logits[i].argmax() }`. + public func verify( + oracleAtHistoryEnd: Int, oracle: (Int) -> Int + ) -> VerifyResult { + if self.token != oracleAtHistoryEnd { + return VerifyResult( + acceptedTokens: [], bonusToken: oracleAtHistoryEnd) + } + let (tokens, parents, _) = flatten() + let n = tokens.count + var childrenOf: [[Int]] = Array(repeating: [], count: n) + for i in 1 ..< n { + childrenOf[parents[i]].append(i) + } + var acceptedPath: [Int] = [tokens[0]] + var currentFlat = 0 + while true { + let predictedNext = oracle(currentFlat) + if let nextFlat = childrenOf[currentFlat] + .first(where: { tokens[$0] == predictedNext }) + { + acceptedPath.append(tokens[nextFlat]) + currentFlat = nextFlat + continue + } + return VerifyResult( + acceptedTokens: acceptedPath, bonusToken: predictedNext) + } + } + + /// Tree-causal additive attention mask for the in-tree positions. + /// Returns a flat `[T·T]` array where `mask[i·T + j]` is: + /// - `0.0` if flat-index `j` is an ancestor of `i` in the tree, + /// or `j == i` itself (the diagonal — every token attends to + /// itself). + /// - `-Float.infinity` otherwise (siblings / cousins / disjoint + /// branches — must NOT attend across alternative paths). + /// + /// Caller adds this mask onto the attention scores BEFORE softmax. + /// The cached-prefix portion (positions < `baseKV`) is always + /// attended (full causal-to-cache) and is NOT included here — + /// wrap this mask into the kernel's `mask` param only for the + /// in-block region. + public func treeCausalMask() -> (mask: [Float], t: Int) { + let (_, parent, _) = flatten() + let t = parent.count + var mask = [Float](repeating: -Float.infinity, count: t * t) + for i in 0 ..< t { + var node = i + while node != -1 { + mask[i * t + node] = 0.0 + node = parent[node] + } + } + return (mask, t) + } +} + +/// A drafter that proposes a tree of candidate continuations. The +/// SpecDecode driver walks the tree, flattens to a tree-causal +/// attention mask, and verifies all candidates in ONE forward pass — +/// accepting the longest matching prefix. +public protocol TreeDrafter: AnyObject { + /// Propose a tree rooted at the first token after `history`. + /// Returns `nil` when the drafter has no confident proposal — the + /// driver falls back to a plain decode step. + /// + /// `maxDepth` caps the depth of the returned tree (root counts as + /// depth 1). `maxNodes` caps the total node count so the driver + /// can budget the verify forward pass. + func proposeTree( + history: [Int], maxDepth: Int, maxNodes: Int + ) -> DraftTreeNode? +} + +extension Drafter { + /// Convert a linear `propose` result into a degenerate tree (one + /// child per node). Used to expose any existing `Drafter` as a + /// `TreeDrafter` without adding a real branching policy. + public func proposeTreeLinear( + history: [Int], maxDepth: Int + ) -> DraftTreeNode? { + let linear = propose(history: history, gamma: maxDepth) + guard let last = linear.last else { return nil } + // Build from leaf inward so children-of-children compose + // correctly. + var node = DraftTreeNode(token: last, children: []) + for t in linear.dropLast().reversed() { + node = DraftTreeNode(token: t, children: [node]) + } + return node + } +} + +/// Adapter: any linear `Drafter` becomes a `TreeDrafter` that emits a +/// degenerate (single-branch) tree. Useful for A/B-ing the +/// tree-verify driver against the linear baseline using the existing +/// `NGramDrafter`. +public final class LinearTreeAdapter: TreeDrafter { + public let inner: Drafter + public init(_ inner: Drafter) { self.inner = inner } + public func proposeTree( + history: [Int], maxDepth: Int, maxNodes _: Int + ) -> DraftTreeNode? { + inner.proposeTreeLinear(history: history, maxDepth: maxDepth) + } +} + +/// Real branching n-gram drafter. +/// +/// Extends `NGramDrafter`'s "scan history for n-gram matches" lookup +/// from "first match → linear chain" to "all matches → top-K +/// continuations per depth → branching tree." Each node expands its +/// `branchingFactor` most-frequent continuations from past +/// occurrences of the current key, recursively up to `maxDepth` or +/// `maxNodes` total. +/// +/// Conforms to both `Drafter` (linear γ fallback via `propose`) and +/// `TreeDrafter` (real branching via `proposeTree`). +public final class NGramTreeDrafter: Drafter, TreeDrafter { + public let maxNMatch: Int + public let minNMatch: Int + /// Number of children per node (top-K continuations). + public let branchingFactor: Int + + public init( + maxNMatch: Int = 3, minNMatch: Int = 2, + branchingFactor: Int = 2 + ) { + precondition( + maxNMatch >= minNMatch && minNMatch >= 1, + "NGramTreeDrafter: maxNMatch (\(maxNMatch)) must be ≥ minNMatch (\(minNMatch)) ≥ 1") + precondition( + branchingFactor >= 1, + "NGramTreeDrafter: branchingFactor must be ≥ 1") + self.maxNMatch = maxNMatch + self.minNMatch = minNMatch + self.branchingFactor = branchingFactor + } + + // MARK: Drafter — linear γ fallback (top-1 chain). + + public func propose(history: [Int], gamma: Int) -> [Int] { + precondition( + gamma >= 0, "NGramTreeDrafter.propose: gamma must be ≥ 0") + guard gamma > 0, !history.isEmpty else { return [] } + var chain: [Int] = [] + chain.reserveCapacity(gamma) + var extended = history + for _ in 0 ..< gamma { + guard let token = topKContinuations(of: extended, k: 1).first + else { break } + chain.append(token) + extended.append(token) + } + return chain + } + + // MARK: TreeDrafter — branching tree. + + public func proposeTree( + history: [Int], maxDepth: Int, maxNodes: Int + ) -> DraftTreeNode? { + guard maxDepth > 0, maxNodes > 0, !history.isEmpty else { + return nil + } + var nodeBudget = maxNodes + let roots = topKContinuations(of: history, k: branchingFactor) + guard let rootTok = roots.first else { return nil } + nodeBudget -= 1 + // Each `roots[i]` becomes a child of the root at depth 1, then + // recurse downwards `maxDepth - 1`. + var rootChildren: [DraftTreeNode] = [] + rootChildren.reserveCapacity(roots.count) + for tok in roots { + guard nodeBudget > 0 else { break } + nodeBudget -= 1 // RESERVE this child's slot before recursing. + var path = history + path.append(tok) + let subtree = buildSubtree( + extendedHistory: path, + remainingDepth: maxDepth - 1, + nodeBudget: &nodeBudget) + rootChildren.append( + DraftTreeNode( + token: tok, children: subtree?.children ?? [])) + } + return DraftTreeNode(token: rootTok, children: rootChildren) + } + + /// Recursive helper: build a chain-or-tree rooted at the implicit + /// next-token, capped by `remainingDepth` and `nodeBudget`. Budget + /// is reserved BEFORE recursing so the cap is a hard upper bound + /// on `tree.size`. + private func buildSubtree( + extendedHistory: [Int], + remainingDepth: Int, + nodeBudget: inout Int + ) -> DraftTreeNode? { + guard remainingDepth > 0, nodeBudget > 0 else { return nil } + let toks = topKContinuations(of: extendedHistory, k: branchingFactor) + guard let firstTok = toks.first else { return nil } + var children: [DraftTreeNode] = [] + children.reserveCapacity(toks.count) + for tok in toks { + guard nodeBudget > 0 else { break } + nodeBudget -= 1 // RESERVE before recursing. + var deeper = extendedHistory + deeper.append(tok) + let sub = buildSubtree( + extendedHistory: deeper, + remainingDepth: remainingDepth - 1, + nodeBudget: &nodeBudget) + children.append( + DraftTreeNode(token: tok, children: sub?.children ?? [])) + } + return DraftTreeNode(token: firstTok, children: children) + } + + /// Top-K most frequent continuations after `history` (longest + /// available n-gram → fall back to shorter). Returns up to `k` + /// tokens sorted by occurrence count descending; ties broken by + /// token id ascending for determinism. + private func topKContinuations(of history: [Int], k: Int) -> [Int] { + guard !history.isEmpty, k > 0 else { return [] } + 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]) + var counts: [Int: Int] = [:] + var probe = keyStart - 1 + while probe >= nMatch - 1 { + if matches(history, at: probe - (nMatch - 1), key: key) { + let candidateIdx = probe + 1 + if candidateIdx < history.count { + counts[history[candidateIdx], default: 0] += 1 + } + } + probe -= 1 + } + if !counts.isEmpty { + return counts.sorted { + if $0.value != $1.value { return $0.value > $1.value } + return $0.key < $1.key + } + .prefix(k) + .map(\.key) + } + } + 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 + } +} diff --git a/Sources/FFAI/Generation/SpecDecode.swift b/Sources/FFAI/Generation/SpecDecode.swift new file mode 100644 index 00000000..be3d5a23 --- /dev/null +++ b/Sources/FFAI/Generation/SpecDecode.swift @@ -0,0 +1,252 @@ +// Copyright 2026 Tom Turney (@TheTom) +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SpecDecode — greedy speculative-decode driver for Qwen3.5/3.6. +// +// One iteration at γ candidates: +// 1. Drafter proposes up to γ candidate next tokens c_0..c_{γ-1}. +// 2. If γ' = drafter's actual proposal length is 0, fall back to a +// regular single-token forward. +// 3. Otherwise: +// a. Snapshot every layer cache. +// b. Run `forwardManyAllLogits([lastAccepted, c_0..c_{γ'-1}], +// startPos=pos)` — this advances every cache by γ' + 1 tokens +// and returns logits at each of the γ' + 1 input positions. +// c. Greedy verify: for i in 0..γ'-1, check whether +// `argmax(logits[i]) == c_i`. Accept consecutive matches up +// to the first mismatch. +// d. If first mismatch at index k: +// * Restore caches from snapshot. +// * Re-run single-step `forward(...)` over `[prev, c_0..c_{k-1}]` +// to advance caches by exactly k+1 tokens (with the +// bit-identical-to-baseline path). +// * Take the next-prev token from the LAST single-step +// forward — `forwardManyAllLogits` can drift slightly +// (different kernel paths / accumulation order) vs the +// single-step decode, so the sampling source for `prev` +// must be the single-step forward output to keep the +// greedy stream bit-identical to the baseline loop. +// * tokens committed this iter: k + 1. +// If no mismatch (all γ' accepted): +// * Commit [c_0..c_{γ'-1}] + the bonus token from the +// trailing logits[γ'] (the model's prediction for what +// comes AFTER the last candidate). +// * tokens committed: γ' + 1. +// * Cache state is already correct — discard snapshot. +// +// Greedy-only for v0. Temperature sampling / nucleus / top-K extensions +// follow the same shape but compare sampled tokens against the +// drafter's proposed tokens via probability-ratio rejection (see the +// original spec-decode paper). + +import Foundation +import Metal + +public struct SpecDecodeStats { + public let tokensGenerated: Int + public let stepsRun: Int + public let candidatesProposed: Int + public let candidatesAccepted: Int + public let fallbackSingleSteps: Int + public let wallclockSeconds: Double + + public var acceptanceRate: Double { + guard candidatesProposed > 0 else { return 0 } + return Double(candidatesAccepted) / Double(candidatesProposed) + } + public var tps: Double { + guard wallclockSeconds > 0 else { return 0 } + return Double(tokensGenerated) / wallclockSeconds + } +} + +public enum SpecDecode { + /// Run a greedy speculative-decode loop. The caller is responsible + /// for prefill — `caches` must already reflect `prompt`, `lastToken` + /// is the most recent (already-sampled) token, `position` is the + /// position where `lastToken` would be inserted on the next decode + /// step (i.e. promptTokens.count if you just finished prefill). + /// + /// Greedy-only — `argmax` everywhere. Stops when `maxNewTokens` is + /// reached or a token in `stopTokens` is emitted. + public static func generateGreedy( + model: Qwen35Model, + drafter: Drafter, + gamma: Int, + lastToken: Int, + position: Int, + caches: [any LayerCacheProtocol], + history: inout [Int], + maxNewTokens: Int, + stopTokens: Set = [], + device: Device = .shared + ) -> SpecDecodeStats { + precondition(gamma >= 1, "SpecDecode.generateGreedy: gamma must be ≥ 1") + precondition( + maxNewTokens >= 0, + "SpecDecode.generateGreedy: maxNewTokens must be ≥ 0") + + var pos = position + var prev = lastToken + var tokensGenerated = 0 + var stepsRun = 0 + var candidatesProposed = 0 + var candidatesAccepted = 0 + var fallbackSingleSteps = 0 + let t0 = Date() + + loop: while tokensGenerated < maxNewTokens { + // ── Draft proposal ──────────────────────────────────────── + let candidates = drafter.propose(history: history + [prev], gamma: gamma) + if candidates.isEmpty { + // No proposal → fall back to single-token decode. + let cmd = device.makeCommandBuffer() + let logits = model.forward( + tokenId: prev, position: pos, + caches: caches, on: cmd, device: device) + cmd.commit() + cmd.waitUntilCompleted() + let next = argmax(logits) + history.append(prev) + pos += 1 + prev = next + tokensGenerated += 1 + stepsRun += 1 + fallbackSingleSteps += 1 + if stopTokens.contains(next) { break loop } + continue + } + let gp = candidates.count + candidatesProposed += gp + + // ── Snapshot caches before the speculative forward ──────── + let snap = caches.snapshotAll(device: device) + + // ── Verify: run forwardManyAllLogits([prev, c_0..c_{gp-1}]) ─ + let inputIds = [prev] + candidates + let cmd = device.makeCommandBuffer() + let allLogits = model.forwardManyAllLogits( + tokenIds: inputIds, startPosition: pos, + caches: caches, on: cmd, device: device) + cmd.commit() + cmd.waitUntilCompleted() + stepsRun += 1 + + // allLogits shape [gp + 1, vocab]. Read host-side argmax per row. + let flat = allLogits.toFloatArray() + let vocab = flat.count / (gp + 1) + precondition( + flat.count == (gp + 1) * vocab, + "SpecDecode: allLogits has \(flat.count) elements; expected (gp+1)*vocab = \((gp + 1) * vocab)" + ) + + // ── Greedy accept loop ──────────────────────────────────── + var acceptedCount = 0 + var firstMismatchTok: Int? = nil + for i in 0 ..< gp { + let rowStart = i * vocab + let modelChoice = argmax(flat, offset: rowStart, count: vocab) + if modelChoice == candidates[i] { + acceptedCount += 1 + } else { + firstMismatchTok = modelChoice + break + } + } + + if firstMismatchTok != nil { + // ── Partial accept (could be 0): restore + replay ──── + // We DO NOT trust the batched verify's logits for + // sampling the next-prev — forwardManyAllLogits can + // drift slightly (different kernel paths, accumulation + // order) vs single forward(). After replay, take the + // next-prev from the LAST single-step forward — that is + // bit-identical to what the baseline greedy loop + // produces. + caches.restoreAll(from: snap, device: device) + let toReplay = [prev] + Array(candidates.prefix(acceptedCount)) + var lastLogits: Tensor! + for (i, tok) in toReplay.enumerated() { + let stepCmd = device.makeCommandBuffer() + lastLogits = model.forward( + tokenId: tok, position: pos + i, + caches: caches, on: stepCmd, device: device) + stepCmd.commit() + stepCmd.waitUntilCompleted() + } + let nextProvenTok = argmax(lastLogits) + history.append(prev) + for c in candidates.prefix(acceptedCount) { history.append(c) } + let totalCommitted = 1 + acceptedCount + pos += totalCommitted + prev = nextProvenTok + tokensGenerated += totalCommitted + candidatesAccepted += acceptedCount + if stopTokens.contains(nextProvenTok) { break loop } + } else { + // ── Full accept: commit all γ' candidates + bonus ──── + // Bonus token = argmax(logits[gp]) — model's prediction + // at position pos + gp + 1. + let bonusStart = gp * vocab + let bonusTok = argmax(flat, offset: bonusStart, count: vocab) + history.append(prev) + for c in candidates { history.append(c) } + let totalCommitted = 1 + gp + pos += totalCommitted + prev = bonusTok + tokensGenerated += totalCommitted + candidatesAccepted += gp + if stopTokens.contains(bonusTok) { break loop } + } + } + // `prev` is the next-iteration's input (not yet a generated + // token — its KV slot isn't in the cache). Caller continues + // from here by calling again with this `prev`, OR can append it + // themselves if they're done generating. We do NOT append it + // here — `tokensGenerated` and `history` reflect ONLY tokens + // committed to the cache, matching the baseline greedy loop's + // semantics. + + return SpecDecodeStats( + tokensGenerated: tokensGenerated, + stepsRun: stepsRun, + candidatesProposed: candidatesProposed, + candidatesAccepted: candidatesAccepted, + fallbackSingleSteps: fallbackSingleSteps, + wallclockSeconds: Date().timeIntervalSince(t0)) + } +} + +// ─── argmax helpers ────────────────────────────────────────────────── + +@inline(__always) +private func argmax(_ logits: Tensor) -> Int { + let host = logits.toFloatArray() + return argmax(host, offset: 0, count: host.count) +} + +@inline(__always) +private func argmax(_ flat: [Float], offset: Int, count: Int) -> Int { + precondition(count > 0) + var bestIdx = 0 + var bestVal = flat[offset] + for i in 1 ..< count { + let v = flat[offset + i] + if v > bestVal { + bestVal = v + bestIdx = i + } + } + return bestIdx +} diff --git a/Sources/FFAI/KVCache/CacheSnapshot.swift b/Sources/FFAI/KVCache/CacheSnapshot.swift new file mode 100644 index 00000000..d9c74011 --- /dev/null +++ b/Sources/FFAI/KVCache/CacheSnapshot.swift @@ -0,0 +1,110 @@ +// Copyright 2026 Tom Turney (@TheTom) +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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) +} + +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). + public 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. + public 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) + composite.setLength(length) + case let (gdn as GDNStateCache, .gdn(currentT, length)): + gdn.restore(from: currentT, device: device) + gdn.setLength(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)" + ) + } + } + } +} diff --git a/Sources/FFAI/KVCache/ConvStateCache.swift b/Sources/FFAI/KVCache/ConvStateCache.swift index b9cbfccc..8d691b84 100644 --- a/Sources/FFAI/KVCache/ConvStateCache.swift +++ b/Sources/FFAI/KVCache/ConvStateCache.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -48,6 +48,10 @@ public final class ConvStateCache: @unchecked Sendable { shape: [kernelSize - 1, nChannels], dtype: dtype, device: device) self.state.zero() + // Conv rolling window persists across every decode step; pin it + // in the residency set so the driver skips per-dispatch + // residency validation on the read+shift that fires each token. + device.markWeightsResident([self.state.buffer]) } /// Reset to zero. Used between sessions; cheap (small fp16/bf16 buffer). @@ -58,6 +62,66 @@ public final class ConvStateCache: @unchecked Sendable { public var bytesAllocated: Int { (kernelSize - 1) * nChannels * dtype.byteSize } + + /// Cached snapshot tensor. Reused across snapshot calls so + /// spec-decode doesn't churn ~30 KB per layer per verify step. + private var cachedSnapshot: Tensor? + + /// Copy `state` into a fresh (or cached) snapshot tensor. Used by + /// spec-decode to roll back the conv rolling window on draft reject. + /// + /// Reuse contract: this method returns a single per-instance + /// scratch tensor. Calling `snapshot()` a second time before + /// `restore(from:)` will overwrite the prior snapshot. Nested or + /// concurrent snapshot usage on the same cache is not supported. + public func snapshot(device: Device = .shared) -> Tensor { + let shape = [kernelSize - 1, nChannels] + if cachedSnapshot == nil + || cachedSnapshot!.shape != shape + || cachedSnapshot!.dtype != dtype + { + cachedSnapshot = Tensor.empty( + shape: shape, dtype: dtype, device: device) + } + let snap = cachedSnapshot! + let cmd = device.makeCommandBuffer() + guard let blit = cmd.makeBlitCommandEncoder() else { + preconditionFailure( + "ConvStateCache.snapshot: makeBlitCommandEncoder failed") + } + let bytes = bytesAllocated + 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 `state` from a snapshot taken via `snapshot()`. + public func restore(from snapshot: Tensor, device: Device = .shared) { + precondition( + snapshot.elementCount == state.elementCount, + "ConvStateCache.restore: snapshot element count mismatch") + precondition( + snapshot.dtype == dtype, + "ConvStateCache.restore: snapshot dtype mismatch") + let cmd = device.makeCommandBuffer() + guard let blit = cmd.makeBlitCommandEncoder() else { + preconditionFailure( + "ConvStateCache.restore: makeBlitCommandEncoder failed") + } + let bytes = bytesAllocated + blit.copy( + from: snapshot.buffer, sourceOffset: snapshot.offset, + to: state.buffer, destinationOffset: state.offset, + size: bytes) + blit.endEncoding() + cmd.commit() + cmd.waitUntilCompleted() + } } extension Array where Element == ConvStateCache { diff --git a/Sources/FFAI/KVCache/GDNStateCache.swift b/Sources/FFAI/KVCache/GDNStateCache.swift index f0162591..e45dd50a 100644 --- a/Sources/FFAI/KVCache/GDNStateCache.swift +++ b/Sources/FFAI/KVCache/GDNStateCache.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -99,6 +99,11 @@ public final class GDNStateCache: LayerCacheProtocol, @unchecked Sendable { self.next = Tensor.empty(shape: shape, dtype: .f32, device: device) self.current.zero() self.next.zero() + // Recurrent state buffers persist across every decode step; pin + // them in the residency set so the driver skips per-dispatch + // residency validation on the read/write/swap that fires each + // token through the GDN layer. + device.markWeightsResident([self.current.buffer, self.next.buffer]) } /// Exchange `current` and `next`. Call this after `Ops.gatedDeltaStep` @@ -129,6 +134,79 @@ public final class GDNStateCache: LayerCacheProtocol, @unchecked Sendable { public var bytesInUse: Int { length == 0 ? 0 : bytesAllocated } + + /// Cached snapshot tensor. Reused across snapshot calls so + /// spec-decode doesn't churn ~2 MB per layer per verify step. + private var cachedSnapshot: Tensor? + + /// Copy `current` into a fresh (or cached) snapshot tensor. Used + /// by spec-decode to roll back the recurrent state on draft reject. + /// Returns the snapshot tensor; caller stores it until needed. + /// + /// Reuse contract: this method returns a single per-instance + /// scratch tensor. Calling `snapshot()` a second time before + /// `restore(from:)` will overwrite the prior snapshot. Nested or + /// concurrent snapshot usage on the same cache is not supported. + 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! + 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 must `setLength(...)` to match. + 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") + 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() + } + + /// 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 + } } extension Array where Element == GDNStateCache { diff --git a/Sources/FFAI/KVCache/KVCache.swift b/Sources/FFAI/KVCache/KVCache.swift index 43e46a25..428ea9f2 100644 --- a/Sources/FFAI/KVCache/KVCache.swift +++ b/Sources/FFAI/KVCache/KVCache.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -195,6 +195,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 entire 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 @@ -280,6 +284,49 @@ public final class KVCache: KVCacheProtocol, @unchecked Sendable { } } + /// Batched range append: takes contiguous `[T, nKVHeads, headDim]` + /// flat K + V tensors and writes them all in ONE shared encoder + /// (2 dispatches: K then V) using a `[T]` u32 positions buffer. + /// Caller provides the positions buffer (allocated + filled with + /// the freshly-reserved slot indices). Replaces the T-loop of + /// `Ops.kvCacheUpdateKV` calls. + public func appendRangeOnGPUMany( + kFlat: Tensor, vFlat: Tensor, + t: Int, positions: Tensor, + on cmd: MTLCommandBuffer + ) { + precondition( + kFlat.dtype == dtype && vFlat.dtype == dtype, + "KVCache.appendRangeOnGPUMany: dtype mismatch") + precondition( + positions.dtype == .u32 && positions.elementCount == t, + "KVCache.appendRangeOnGPUMany: positions must be .u32[T]") + Ops.kvCacheUpdateKVMany( + kSrc: kFlat, kCache: kBuffer, + vSrc: vFlat, vCache: vBuffer, + positions: positions, t: t, + nKVHeads: nKVHeads, headDim: headDim, maxSeq: maxSeq, + on: cmd) + } + + /// Reserve T sequential physical slots in the cache, writing the + /// chosen indices into `positionsOut` (a u32 buffer length ≥ T). + /// Atomic under `lengthLock`. Returns nothing — caller passes the + /// positions tensor straight to `appendRangeOnGPUMany`. + public func reserveSlotsManyOnHost(t: Int, into positionsOut: Tensor) { + precondition( + positionsOut.dtype == .u32, + "KVCache.reserveSlotsManyOnHost: positionsOut must be .u32") + precondition( + positionsOut.elementCount >= t, + "KVCache.reserveSlotsManyOnHost: positionsOut shorter than T") + let ptr = positionsOut.buffer.contents().advanced(by: positionsOut.offset) + .bindMemory(to: UInt32.self, capacity: t) + lengthLock.withLock { + for r in 0 ..< t { ptr[r] = UInt32(_evictionState.reserveNextSlot()) } + } + } + /// Write one timestep's K/V at an explicit physical slot **without** /// touching `length`. Diffusion-block forwards stage their scratch /// K/V in the buffer's free region `[length, maxSeq)` across denoise diff --git a/Sources/FFAI/Layers.swift b/Sources/FFAI/Layers.swift index 7873e3b2..08dfd792 100644 --- a/Sources/FFAI/Layers.swift +++ b/Sources/FFAI/Layers.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -221,15 +221,16 @@ public final class QuantizedLinear: Module { ) -> 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. + // 4-bit goes through the fast mt_qmm_mma path when the output + // dimension is a multiple of 32 (the BN tile of the underlying + // `dequantGemmDynamicM` kernel). Anything narrower — including + // 4-bit shared_expert_gate at outDim=1 — falls through to the + // per-row dequantGemv loop. Other bit-widths (commonly 8-bit + // on the same narrow projections) take the same fallback. Both + // fallbacks are bit-identical to the per-token path; their + // per-token launch overhead is in the noise on hidden→1 shapes. let out: Tensor - if bits == 4 { + if bits == 4 && outDim % 32 == 0 { out = Tensor.empty(shape: [t, outDim], dtype: x.dtype, device: device) Ops.dequantGemmDynamicM( input: x, weight: weight, scales: scales, biases: biases, diff --git a/Sources/FFAI/Loader/Model.swift b/Sources/FFAI/Loader/Model.swift index f86d6c94..9764d224 100644 --- a/Sources/FFAI/Loader/Model.swift +++ b/Sources/FFAI/Loader/Model.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -18,6 +18,7 @@ // surface. import Foundation +import Metal import Tokenizers public enum ModelError: Error, CustomStringConvertible { @@ -1140,6 +1141,18 @@ public final class Model: @unchecked Sendable { "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 every weight buffer in a persistent MTLResidencySet + // so prefill / decode dispatches skip per-allocation + // residency tracking. macOS 15+ / iOS 18+; older OSes + // and `FFAI_NO_RESIDENCY_SET=1` no-op via + // `Device.markWeightsResident`. + 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/Audio/VAD/FireRedVAD.swift b/Sources/FFAI/Models/Audio/VAD/FireRedVAD.swift index c8c5d7a9..d8b7694d 100644 --- a/Sources/FFAI/Models/Audio/VAD/FireRedVAD.swift +++ b/Sources/FFAI/Models/Audio/VAD/FireRedVAD.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -1199,11 +1199,13 @@ public final class FireRedVADModel: @unchecked Sendable { case 0x8C: // SHORT_BINUNICODE len1 stack.append(.str(readLen1String())) case 0x8D: // BINUNICODE8 - let n = - Int(pkl[pos]) | (Int(pkl[pos + 1]) << 8) | (Int(pkl[pos + 2]) << 16) - | (Int(pkl[pos + 3]) << 24) - | (Int(pkl[pos + 4]) << 32) | (Int(pkl[pos + 5]) << 40) + let nLo32 = + Int(pkl[pos]) | (Int(pkl[pos + 1]) << 8) + | (Int(pkl[pos + 2]) << 16) | (Int(pkl[pos + 3]) << 24) + let nHi32 = + (Int(pkl[pos + 4]) << 32) | (Int(pkl[pos + 5]) << 40) | (Int(pkl[pos + 6]) << 48) | (Int(pkl[pos + 7]) << 56) + let n = nLo32 | nHi32 pos += 8 let s = String(bytes: pkl[pos ..< (pos + n)], encoding: .utf8) ?? "" pos += n diff --git a/Sources/FFAI/Models/MoELayer.swift b/Sources/FFAI/Models/MoELayer.swift index ce8875ce..43a71f76 100644 --- a/Sources/FFAI/Models/MoELayer.swift +++ b/Sources/FFAI/Models/MoELayer.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -307,6 +307,15 @@ public struct MoERouter: Sendable { /// `decode` runs: gate gemv → CPU route → per-expert SwiGLU on the /// selected experts → combine. See the command-buffer contract note in /// the file header — `decode` commits the passed `cmd`. +/// +/// Threading contract: instances are NOT safe for concurrent calls. +/// `decode` / `decodeMany` mutate per-instance scratch buffers +/// (`routerIndicesScratch`, `accumulatorScratch`, `expertGScratches`, +/// `dequantizedGateCache`, …) without internal synchronisation. Callers +/// must serialise inference on a given instance — typical model loops +/// already do this since a token can't fan out across MoE layers in +/// parallel. The class is `final` for closed-world dispatch, not for +/// concurrency. public final class MoELayer: Module, DecoderLayer { /// Router projection: hidden → nExperts logits. public let gate: AnyLinear @@ -390,6 +399,34 @@ public final class MoELayer: Module, DecoderLayer { public let enableBGEMM: Bool public let useBm8Env: Bool public let useM1Env: Bool + /// Additional env-flag caches. + let noDequantGate: Bool + let gpuRouterEnabled: Bool + let forceBGEMMAtT1: Bool + let noBGEMMBm64: Bool + + /// Lazy dequant of the gate weight + cached output. Skipped when + /// `gate.inner` is not a 8-bit QuantizedLinear or env + /// `FFAI_MOE_NO_DEQUANT_GATE=1` is set. + private var dequantizedGateCache: Tensor? + private var dequantizedGateCacheKN: Tensor? + private var dequantizedGateAttempted = false + private var gateLogitsScratch: Tensor? + + /// Cached scratches for the per-expert weighted-sum loop at + /// decode T=1. + private var accumulatorScratch: Tensor? + private var topKScalarsBuf: MTLBuffer? + /// Cached per-expert g/u/inner/down outputs. + private var expertGScratches: [Tensor] = [] + private var expertUScratches: [Tensor] = [] + private var expertInnerScratches: [Tensor] = [] + private var expertOutScratches: [Tensor] = [] + /// GPU MoE router scratches — written once per layer per token by + /// `mt_moe_router_topk`, then consumed by the per-slot indexed qmms + /// and the scalarFMAChain8. + private var routerIndicesScratch: Tensor? + private var routerWeightsScratch: Tensor? /// - gate: hidden → nExperts router projection. /// - gateProj/upProj/downProj: `nExperts`-long arrays of per-expert @@ -450,6 +487,14 @@ public final class MoELayer: Module, DecoderLayer { self.enableBGEMM = env["FFAI_MOE_BGEMM"] != nil self.useBm8Env = env["FFAI_MOE_BGEMM_BM8"] != nil self.useM1Env = env["FFAI_MOE_M1"] != nil + self.noDequantGate = env["FFAI_MOE_NO_DEQUANT_GATE"] != nil + // Default GPU router ON — end-to-end correctness is pinned + // across bf16/f16/f32 by the integration suite. Opt out with + // `FFAI_MOE_GPU_ROUTER=0`. The win: 40 × + // `waitUntilCompleted` per decode token → 0. + self.gpuRouterEnabled = env["FFAI_MOE_GPU_ROUTER"] != "0" + self.forceBGEMMAtT1 = env["FFAI_MOE_BGEMM_FORCE_T1"] != nil + self.noBGEMMBm64 = env["FFAI_MOE_BGEMM_NO_BM64"] != nil } public func parameters() -> [(String, Tensor)] { @@ -492,9 +537,30 @@ 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) + // ── 2a. GPU MoE router fast path ────────────────────────────── + // When `FFAI_MOE_GPU_ROUTER=1` (default ON) AND the layer has + // stacked int4 expert layout AND routing matches + // `mt_moe_router_topk`'s assumptions (softmax-then-topK + Qwen- + // MoE-style renorm, topK=8), skip the CPU sync entirely. The + // router writes topK indices + weights to a GPU buffer; per- + // expert qmms read their slot's expert id from that buffer via + // `dequantGemvInt4ExpertIndexed`; the per-slot routing weight + // feeds `scalarFMAChain8`. + if self.gpuRouterEnabled, + let stacked = stackedInt4Experts, + stacked.dtype == h.dtype, + router.gatingMode == .softmaxThenTopK, + router.normTopKProb, + router.expertBias == nil, + router.topK == 8 + { + return decodeGPURouter( + h: h, logitsTensor: logitsTensor, + stacked: stacked, cmd: cmd, device: device) + } + // ── 2. Commit + wait so the router can read the logits ─────── // The decode path batches a token onto one command buffer; the // router needs a CPU sync point. Commit here, then run the @@ -624,20 +690,56 @@ public final class MoELayer: Module, DecoderLayer { let gateLogitsAll = gate.callMany(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 ───────────────────────────────── + // ── 2. Routing — GPU fast path or CPU sync ────────────────────── + // When router config matches `mt_moe_router_topk`'s semantic + // envelope (softmax → topK + Qwen-MoE renorm, no expert bias, + // k=8), run the T-batched router on GPU via + // `Ops.moeRouterTopKMany`. Eliminates the T-parallel host + // `router.route` loop (~50 ms of CPU work per prefill at T=512 × + // 40 MoE layers). One commit+wait still needed to read indices + + // weights for the host-side sort plan. let nExperts = router.nExperts let topK = router.topK - let logitsHost = gateLogitsAll.toFloatArray() // [T·nExperts] - var routings: [MoERouter.Routing] = [] - routings.reserveCapacity(t) - for r in 0 ..< t { - let start = r * nExperts - let rowLogits = Array(logitsHost[start ..< (start + nExperts)]) - routings.append(router.route(logits: rowLogits)) + var routings = [MoERouter.Routing]( + repeating: MoERouter.Routing(indices: [], weights: []), + count: t) + let useGPURouter = + self.gpuRouterEnabled + && router.gatingMode == .softmaxThenTopK + && router.normTopKProb + && router.expertBias == nil + && router.topK == 8 + if useGPURouter { + let indicesBuf = Tensor.empty( + shape: [t, topK], dtype: .u32, device: device) + let weightsBuf = Tensor.empty( + shape: [t, topK], dtype: gateLogitsAll.dtype, device: device) + Ops.moeRouterTopKMany( + logits: gateLogitsAll, indicesOut: indicesBuf, + weightsOut: weightsBuf, + t: t, nExperts: nExperts, k: topK, + normTopkProb: router.normTopKProb, on: cmd) + cmd.commit() + cmd.waitUntilCompleted() + let indicesHost: [UInt32] = indicesBuf.toArray(as: UInt32.self) + let weightsHost: [Float] = weightsBuf.toFloatArray() + for r in 0 ..< t { + let base = r * topK + let idxs = (0 ..< topK).map { Int(indicesHost[base + $0]) } + let wts = Array(weightsHost[base ..< (base + topK)]) + routings[r] = MoERouter.Routing(indices: idxs, weights: wts) + } + } else { + cmd.commit() + cmd.waitUntilCompleted() + let logitsHost = gateLogitsAll.toFloatArray() + 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 ──────────────────────── @@ -756,9 +858,10 @@ public final class MoELayer: Module, DecoderLayer { // 2.69× T=32 win. // - mTotal ≤ 8 + `FFAI_MOE_BGEMM_BM8=1` → bm8 (decode T=1 // fallback). - let useBm64 = - mTotal >= 64 - && ProcessInfo.processInfo.environment["FFAI_MOE_BGEMM_BM64"] != nil + // bm64_mpp default-on for mTotal ≥ 1024. Crossover sits between + // mTotal=512 (bm16 +19%) and mTotal=1024 (bm64 +18%). Opt out + // via FFAI_MOE_BGEMM_NO_BM64=1. + let useBm64 = mTotal >= 1024 && !self.noBGEMMBm64 let useBm8 = !useBm64 && topK <= 8 && useBm8Env && mTotal <= 8 let bgemm: (Tensor, Tensor, Tensor, Tensor, Tensor, Int, Int, Int, Int, MTLCommandBuffer, Tensor) @@ -1043,6 +1146,192 @@ public final class MoELayer: Module, DecoderLayer { let inner = Ops.swiglu(gate: g, up: u, on: cmd) return downProj(inner, on: cmd) } + + /// GPU MoE router decode path. Replaces the + /// `cmd.commit + waitUntilCompleted` sync with a full-GPU pipeline: + /// router writes topK indices + weights to GPU buffers; per-slot + /// indexed qmms (gate/up phase + down phase) batch onto shared + /// encoders; final accumulation via `scalarFMAChain8` (or the fused + /// `moeDownSwigluAccumInt4Chain8` kernel when constraints allow). + /// + /// Predicated by `decode` callers — the layer must have stackedInt4 + /// experts with matching dtype, router must be softmax-then-topK + /// with normTopKProb=true / no expertBias / topK=8. + private func decodeGPURouter( + h: Tensor, logitsTensor: Tensor, + stacked: StackedInt4Experts, + cmd: MTLCommandBuffer, device: Device + ) -> Tensor { + // Lazy-init per-instance scratches. + if routerIndicesScratch == nil + || routerIndicesScratch!.elementCount != router.topK + { + routerIndicesScratch = Tensor.empty( + shape: [router.topK], dtype: .u32, device: device) + routerWeightsScratch = Tensor.empty( + shape: [router.topK], dtype: h.dtype, device: device) + } + if accumulatorScratch == nil + || accumulatorScratch!.dtype != h.dtype + || accumulatorScratch!.elementCount != hidden + { + accumulatorScratch = Tensor.empty( + shape: [hidden], dtype: h.dtype, device: device) + } + let outDim = stacked.moeIntermediate + for slot in 0 ..< router.topK { + while expertGScratches.count <= slot { + expertGScratches.append( + Tensor.empty(shape: [outDim], dtype: h.dtype, device: device)) + expertUScratches.append( + Tensor.empty(shape: [outDim], dtype: h.dtype, device: device)) + expertInnerScratches.append( + Tensor.empty(shape: [outDim], dtype: h.dtype, device: device)) + expertOutScratches.append( + Tensor.empty(shape: [hidden], dtype: h.dtype, device: device)) + } + } + + // GPU router: logits → topK indices + weights (still on cmd). + Ops.moeRouterTopK( + logits: logitsTensor, + indicesOut: routerIndicesScratch!, + weightsOut: routerWeightsScratch!, + nExperts: router.nExperts, k: router.topK, + normTopkProb: router.normTopKProb, + on: cmd) + + // Batched per-expert indexed qmms on ONE encoder per phase. + // Gate (8 calls) + up (8 calls) share inDim/outDim/groupSize → + // one Many encoder. Down (8 calls) has different out_dim → + // second Many encoder. + var expertIdxScratches: [Tensor] = [] + expertIdxScratches.reserveCapacity(router.topK) + for slot in 0 ..< router.topK { + expertIdxScratches.append( + Tensor( + buffer: routerIndicesScratch!.buffer, + offset: routerIndicesScratch!.offset + slot * 4, + shape: [1], dtype: .u32)) + } + + // Phase 1a: 16 calls (8 gate + 8 up) on one encoder. + var ws: [Tensor] = [] + var ss: [Tensor] = [] + var bs: [Tensor] = [] + var ins: [Tensor] = [] + var idxs: [Tensor] = [] + var outs0: [Tensor] = [] + ws.reserveCapacity(router.topK * 2) + ss.reserveCapacity(router.topK * 2) + bs.reserveCapacity(router.topK * 2) + ins.reserveCapacity(router.topK * 2) + idxs.reserveCapacity(router.topK * 2) + outs0.reserveCapacity(router.topK * 2) + for slot in 0 ..< router.topK { + ws.append(stacked.gateWeight) + ss.append(stacked.gateScales) + bs.append(stacked.gateBiases) + ins.append(h) + idxs.append(expertIdxScratches[slot]) + outs0.append(expertGScratches[slot]) + ws.append(stacked.upWeight) + ss.append(stacked.upScales) + bs.append(stacked.upBiases) + ins.append(h) + idxs.append(expertIdxScratches[slot]) + outs0.append(expertUScratches[slot]) + } + Ops.dequantGemvInt4ExpertIndexedMany( + weightsStacked: ws, scalesStacked: ss, biasesStacked: bs, + inputs: ins, expertIndices: idxs, outputs: outs0, + groupSize: stacked.groupSize, on: cmd) + + // Fuse phase 1b (swigluMany), phase 2 (8 down indexed-gemvs), + // and phase 3 (scalarFMAChain8) into ONE dispatch via + // `ffai_moe_down_swiglu_accum_int4_chain8`. Per-slot `inner` is + // staged in 3 KiB threadgroup memory, never spilling to global + // memory. Kernel constraint: moeIntermediate ≤ 768 + groupSize + // == 64. Qwen3.6-A3B has exactly 768 so this matches; fall back + // to the 3-dispatch chain for any larger configs. + let gs = Array(expertGScratches.prefix(router.topK)) + let us = Array(expertUScratches.prefix(router.topK)) + let acc = accumulatorScratch! + if router.topK == 8 + && stacked.moeIntermediate <= 768 + && stacked.groupSize == 64 + { + Ops.moeDownSwigluAccumInt4Chain8( + gates: gs, ups: us, + expertIndices: routerIndicesScratch!, + slotWeights: routerWeightsScratch!, + weightsStacked: stacked.downWeight, + scalesStacked: stacked.downScales, + biasesStacked: stacked.downBiases, + output: acc, + inDim: stacked.moeIntermediate, + outDim: hidden, + groupSize: stacked.groupSize, + on: cmd) + } else { + // Legacy 3-dispatch chain (Phase 1b + 2 + 3). + let inners = Array(expertInnerScratches.prefix(router.topK)) + Ops.swigluMany(gates: gs, ups: us, outs: inners, on: cmd) + var dws: [Tensor] = [] + var dss: [Tensor] = [] + var dbs: [Tensor] = [] + var dins: [Tensor] = [] + var didxs: [Tensor] = [] + var douts: [Tensor] = [] + dws.reserveCapacity(router.topK) + dss.reserveCapacity(router.topK) + dbs.reserveCapacity(router.topK) + dins.reserveCapacity(router.topK) + didxs.reserveCapacity(router.topK) + douts.reserveCapacity(router.topK) + for slot in 0 ..< router.topK { + dws.append(stacked.downWeight) + dss.append(stacked.downScales) + dbs.append(stacked.downBiases) + dins.append(expertInnerScratches[slot]) + didxs.append(expertIdxScratches[slot]) + douts.append(expertOutScratches[slot]) + } + Ops.dequantGemvInt4ExpertIndexedMany( + weightsStacked: dws, scalesStacked: dss, biasesStacked: dbs, + inputs: dins, expertIndices: didxs, outputs: douts, + groupSize: stacked.groupSize, on: cmd) + var scalars: [Tensor] = [] + scalars.reserveCapacity(router.topK) + for slot in 0 ..< router.topK { + scalars.append( + Tensor( + buffer: routerWeightsScratch!.buffer, + offset: routerWeightsScratch!.offset + + slot * h.dtype.byteSize, + shape: [1], dtype: h.dtype)) + } + let outs = Array(expertOutScratches.prefix(router.topK)) + Ops.scalarFMAChain8( + scalars: scalars, values: outs, out: acc, on: cmd) + } + + // Commit cmd without wait — preserves the caller contract + // (caller starts a fresh cmd for residual-add / shared expert). + // The win is the ABSENT `waitUntilCompleted` that the CPU-sync + // path needed before this branch. + // + // Contract: the returned tensor aliases `accumulatorScratch`, + // which the GPU is still writing through the in-flight `cmd`. + // Callers MUST not re-invoke `decode` on this layer until that + // cmd completes — a second call would re-enter `decodeGPURouter` + // and start overwriting the scratch while the GPU is still + // reading from it. Per-token serial dispatch in the model loop + // satisfies this; concurrent decodes on a shared MoELayer + // instance do not. + cmd.commit() + return acc + } } // ─── Integration note for MoE-bearing families ─────────────────────── diff --git a/Sources/FFAI/Models/Text/Granite4Text.swift b/Sources/FFAI/Models/Text/Granite4Text.swift index 5bda00be..828172b5 100644 --- a/Sources/FFAI/Models/Text/Granite4Text.swift +++ b/Sources/FFAI/Models/Text/Granite4Text.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -673,7 +673,9 @@ public final class Granite4DenseMLP: 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) + // Fused silu(gate) * up — one kernel, intermediate silu(g) + // stays in registers instead of round-tripping to DRAM. + let inner = Ops.swiglu(gate: g, up: u, on: cmd) return downProj(inner, on: cmd) } } diff --git a/Sources/FFAI/Models/Text/JambaText.swift b/Sources/FFAI/Models/Text/JambaText.swift index 8bbcc690..dd560e68 100644 --- a/Sources/FFAI/Models/Text/JambaText.swift +++ b/Sources/FFAI/Models/Text/JambaText.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -633,7 +633,9 @@ 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) + // Fused silu(gate) * up — one kernel, intermediate silu(g) + // stays in registers instead of round-tripping to DRAM. + let inner = Ops.swiglu(gate: g, up: u, on: cmd) return downProj(inner, on: cmd) } } diff --git a/Sources/FFAI/Models/Text/LFM2Text.swift b/Sources/FFAI/Models/Text/LFM2Text.swift index 0855b86f..0409ef65 100644 --- a/Sources/FFAI/Models/Text/LFM2Text.swift +++ b/Sources/FFAI/Models/Text/LFM2Text.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -627,7 +627,9 @@ 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) + // Fused silu(gate) * up — one kernel, intermediate stays in + // registers instead of materialising silu(gate) in DRAM. + let inner = Ops.swiglu(gate: gate, up: up, on: cmd) return w2(inner, on: cmd) } } diff --git a/Sources/FFAI/Models/Text/Qwen3xText.swift b/Sources/FFAI/Models/Text/Qwen3xText.swift index bc9a46c6..9c7bcb8c 100644 --- a/Sources/FFAI/Models/Text/Qwen3xText.swift +++ b/Sources/FFAI/Models/Text/Qwen3xText.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -417,7 +417,7 @@ public struct Qwen35Hybrid: Qwen35Variant { qNorm: qNorm, kNorm: kNorm, nHeads: nHeads, nKVHeads: nKVHeads, headDim: headDim, rotaryDim: rotaryDim, ropeTheta: ropeTheta, - attnOutputGate: attnOutputGate) + attnOutputGate: attnOutputGate, device: device) layers.append( Qwen35AttentionLayer( inputNorm: inputNorm, postNorm: postNorm, @@ -777,17 +777,17 @@ public final class Qwen35DenseMLP: Module { return out } - /// down(silu(gate(x)) * up(x)). + /// down(silu(gate(x)) * up(x)). Uses fused `Ops.swiglu` — + /// one dispatch vs silu + mul = two dispatches. 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) } /// T-batched: down(silu(gate(x)) * up(x)) over T rows. Returns - /// `[T, hidden]` flat. Three batched projections + one elementwise - /// SwiGLU pass. + /// `[T, hidden]` flat. Three batched projections + one fused SwiGLU. func forwardMany( _ xNormFlat: Tensor, t: Int, cmd: MTLCommandBuffer, device: Device @@ -795,7 +795,7 @@ public final class Qwen35DenseMLP: Module { let xRows = xNormFlat.reshaped(to: [t, xNormFlat.elementCount / t]) let g = gateProj.callMany(xRows, t: t, on: cmd, device: device) let u = upProj.callMany(xRows, t: t, on: cmd, device: device) - let inner = Ops.mul(Ops.silu(g, on: cmd), u, on: cmd) + let inner = Ops.swiglu(gate: g, up: u, on: cmd) return downProj.callMany(inner, t: t, on: cmd, device: device) } } @@ -810,6 +810,13 @@ public final class Qwen35DenseMLP: Module { // `MoELayer.decode` commits the command buffer; this wrapper therefore // also commits, and runs the shared expert + the routed-combine add on // fresh private buffers so the returned tensor is fully resident. +// +// Threading contract: instances are NOT safe for concurrent calls. +// `forward` / `forwardMany` mutate per-instance shared-expert scratch +// buffers (`sharedSgScratch`, `sharedSuScratch`, +// `sharedGateLogitScratch`, `sharedResultScratch`) without internal +// synchronisation. Callers must serialise inference on a given +// instance. public final class Qwen35MoEFFN: Module { let moe: MoELayer @@ -817,6 +824,16 @@ public final class Qwen35MoEFFN: Module { let sharedExpertGate: AnyLinear let hidden: Int + /// Cached shared-expert scratches. The single-token forward fans + /// gate+up + expert-gate qmms through one shared encoder; pre- + /// allocating the outputs at instance level avoids per-token + /// Tensor.empty allocations. + private var sharedSgScratch: Tensor? + private var sharedSuScratch: Tensor? + private var sharedGateLogitScratch: Tensor? + /// Cached `result` output for the GPU fan-out. + private var sharedResultScratch: Tensor? + init( moe: MoELayer, sharedGateProj: AnyLinear, sharedUpProj: AnyLinear, @@ -853,37 +870,120 @@ public final class Qwen35MoEFFN: Module { /// Run the MoE FFN. `MoELayer.decode` commits the passed `cmd`; the /// shared expert + the final add run on fresh private buffers, so /// the returned tensor never depends on the now-dead `cmd`. + /// + /// Optional `residual` parameter folds the post-FFN residual add + /// into the final `sigmoidScalarFMA` dispatch (becomes + /// `sigmoidScalarFMAResidual`). When `residual != nil`, the returned + /// tensor already includes `residual + routed + sigmoid(gate)·sharedOut` + /// and the caller MUST NOT do its own `Ops.add(residual, ffnOut)`. func forward( _ xNorm: Tensor, position: Int, - cmd: MTLCommandBuffer, device: Device + cmd: MTLCommandBuffer, device: Device, + residual: Tensor? = nil ) -> Tensor { // Routed top-K experts — commits `cmd`. let routed = moe.decode( xNorm, position: position, cache: StatelessLayerCache(), cmd: cmd, device: device) - // Shared expert on a fresh buffer: SwiGLU + scalar gate logit. - // The fused-sigmoid kernel keeps the scalar on the GPU, so this - // buffer no longer needs a `waitUntilCompleted` — it just needs - // to be in flight so the FMA can hazard-track its dependencies. + // Shared expert on a fresh buffer. let work = device.makeCommandBuffer() - let sg = sharedGateProj(xNorm, on: work) - let su = sharedUpProj(xNorm, on: work) - let sharedInner = Ops.mul(Ops.silu(sg, on: work), su, on: work) + // Batch sharedGate + sharedUp + sharedExpertGate (3 int4 qmms + // sharing the same `xNorm`) onto ONE shared encoder via + // dequantGemvInt4Three. Saves ~40 encoder begin/end pairs per + // decode token on Qwen3.6-A3B (40 layers × one pair each). + let sg: Tensor + let su: Tensor + let gateLogit: Tensor + let qg = sharedGateProj.inner as? QuantizedLinear + let qu = sharedUpProj.inner as? QuantizedLinear + let qgate = sharedExpertGate.inner as? QuantizedLinear + if let qg = qg, let qu = qu, let qgate = qgate, + qg.bits == 4, qu.bits == 4, qgate.bits == 4, + qg.groupSize == qu.groupSize, qu.groupSize == qgate.groupSize + { + let outDim = qg.weight.shape[0] + let gateOutDim = qgate.weight.shape[0] + if sharedSgScratch == nil + || sharedSgScratch!.dtype != xNorm.dtype + || sharedSgScratch!.elementCount != outDim + { + sharedSgScratch = Tensor.empty( + shape: [outDim], dtype: xNorm.dtype, device: device) + sharedSuScratch = Tensor.empty( + shape: [outDim], dtype: xNorm.dtype, device: device) + } + if sharedGateLogitScratch == nil + || sharedGateLogitScratch!.dtype != xNorm.dtype + || sharedGateLogitScratch!.elementCount != gateOutDim + { + sharedGateLogitScratch = Tensor.empty( + shape: [gateOutDim], dtype: xNorm.dtype, device: device) + } + sg = sharedSgScratch! + su = sharedSuScratch! + gateLogit = sharedGateLogitScratch! + Ops.dequantGemvInt4Three( + input: xNorm, + w0: qg.weight, s0: qg.scales, b0: qg.biases, out0: sg, + w1: qu.weight, s1: qu.scales, b1: qu.biases, out1: su, + w2: qgate.weight, s2: qgate.scales, b2: qgate.biases, + out2: gateLogit, + groupSize: qg.groupSize, on: work) + } else if let qg = qg, let qu = qu, + qg.bits == 4, qu.bits == 4, qg.groupSize == qu.groupSize + { + let outDim = qg.weight.shape[0] + if sharedSgScratch == nil + || sharedSgScratch!.dtype != xNorm.dtype + || sharedSgScratch!.elementCount != outDim + { + sharedSgScratch = Tensor.empty( + shape: [outDim], dtype: xNorm.dtype, device: device) + sharedSuScratch = Tensor.empty( + shape: [outDim], dtype: xNorm.dtype, device: device) + } + sg = sharedSgScratch! + su = sharedSuScratch! + Ops.dequantGemvInt4Two( + input: xNorm, + w0: qg.weight, s0: qg.scales, b0: qg.biases, out0: sg, + w1: qu.weight, s1: qu.scales, b1: qu.biases, out1: su, + groupSize: qg.groupSize, on: work) + gateLogit = sharedExpertGate(xNorm, on: work) + } else { + sg = sharedGateProj(xNorm, on: work) + su = sharedUpProj(xNorm, on: work) + gateLogit = sharedExpertGate(xNorm, on: work) + } + // Fused SwiGLU. + let sharedInner = Ops.swiglu(gate: sg, up: su, on: work) let sharedOut = sharedDownProj(sharedInner, on: work) - let gateLogit = sharedExpertGate(xNorm, on: work) work.commit() // GPU fan-out: `out = routed + sigmoid(gateLogit) * sharedOut` - // in one dispatch. Replaces the prior host detour - // (`gateLogit.toFloatArray()` + host sigmoid + `Tensor.filled` - // broadcast + mul + add + commit + wait) — saves one host stall - // per MoE layer per token. Fires on all 40 Qwen3.6-A3B layers. + // (+ residual when the caller passes one) in one dispatch. let fmaCmd = device.makeCommandBuffer() - let result = Tensor.empty(shape: [hidden], dtype: routed.dtype, device: device) - Ops.sigmoidScalarFMA( - gate: gateLogit, value: sharedOut, base: routed, - into: result, on: fmaCmd) + if sharedResultScratch == nil + || sharedResultScratch!.dtype != routed.dtype + || sharedResultScratch!.elementCount != hidden + { + sharedResultScratch = Tensor.empty( + shape: [hidden], dtype: routed.dtype, device: device) + } + let result = sharedResultScratch! + if let residual = residual { + // Fused 4-input. result = residual + routed + + // sigmoid(gateLogit) · sharedOut. Saves 1 dispatch + 1 + // [hidden] DRAM roundtrip vs sigmoidScalarFMA then Ops.add. + Ops.sigmoidScalarFMAResidual( + gate: gateLogit, value: sharedOut, base: routed, + residual: residual, into: result, on: fmaCmd) + } else { + Ops.sigmoidScalarFMA( + gate: gateLogit, value: sharedOut, base: routed, + into: result, on: fmaCmd) + } fmaCmd.commit() return result } @@ -911,7 +1011,8 @@ public final class Qwen35MoEFFN: Module { let xRows = xNormFlat.reshaped(to: [t, hidden]) let sgAll = sharedGateProj.callMany(xRows, t: t, on: work, device: device) let suAll = sharedUpProj.callMany(xRows, t: t, on: work, device: device) - let sharedInnerAll = Ops.mul(Ops.silu(sgAll, on: work), suAll, on: work) + // Fused SwiGLU. + let sharedInnerAll = Ops.swiglu(gate: sgAll, up: suAll, on: work) let sharedOutAll = sharedDownProj.callMany( sharedInnerAll, t: t, on: work, device: device) @@ -1012,6 +1113,17 @@ public final class Qwen35GDNLayerCache: LayerCacheProtocol, @unchecked Sendable public var bytesInUse: Int { length == 0 ? 0 : bytesAllocated } + + /// Composite-level length-only restore. Tensor state for conv/gdn is + /// restored separately by the caller via the per-cache `restore` + /// methods; this fixes the position counter without disturbing the + /// just-restored sub-cache buffers. + public func setLength(_ length: Int) { + precondition( + length >= 0, + "Qwen35GDNLayerCache.setLength: must be ≥ 0") + self.length = length + } } // ─── Qwen35GDNMixer — Gated Delta Net recurrent mixer ──────────────── @@ -1023,6 +1135,12 @@ public final class Qwen35GDNLayerCache: LayerCacheProtocol, @unchecked Sendable // Because the prep needs a CPU sync, `forward` commits the command // buffer it is handed and returns a resident tensor on a fresh buffer // (the Jamba mamba-mixer pattern). +// +// Threading contract: instances are NOT safe for concurrent calls. +// `forward` / `forwardMany` / `forwardManyChunked` mutate per-instance +// scratch buffers and reuse them across calls without internal +// synchronisation. Callers must serialise inference on a given +// instance. public final class Qwen35GDNMixer: Module { let inProjQKV, inProjZ, inProjB, inProjA, outProj: AnyLinear @@ -1086,6 +1204,28 @@ public final class Qwen35GDNMixer: Module { let yF32Scratch: Tensor let yGatedScratch: Tensor + /// Batched 4-projection input fast path. When all 4 inProj are + /// QuantizedLinear int4 with same groupSize, the mixer routes + /// through `Ops.dequantGemvInt4Four` — 4 qmms on one shared + /// encoder. Saves 3 encoder begin/end pairs per GDN layer × 30 ≈ + /// 1.5 ms/token. Allocated scratch buffers are pinned in the + /// device residency set in init. + private var batchedInputProj: + ( + qW: Tensor, qS: Tensor, qB: Tensor, + zW: Tensor, zS: Tensor, zB: Tensor, + bW: Tensor, bS: Tensor, bB: Tensor, + aW: Tensor, aS: Tensor, aB: Tensor, + groupSize: Int + )? + private var qkvScratch: Tensor? + private var zScratch: Tensor? + private var bRawScratch: Tensor? + private var aRawScratch: Tensor? + /// Single-dispatch `ffai_batched_4_qgemv_fast` fast path. + /// Constraints: in_dim%512==0, each out_dim%8==0, groupSize=64. + private var fused4Eligible: Bool = false + init( inProjQKV: AnyLinear, inProjZ: AnyLinear, inProjB: AnyLinear, inProjA: AnyLinear, outProj: AnyLinear, @@ -1133,13 +1273,14 @@ public final class Qwen35GDNMixer: Module { self.aLogTF32 = makeF32Tensor35(aLog, device: device) self.dtBiasTF32 = makeF32Tensor35(dtBias, device: device) self.epsBufFused = makeF32Tensor35([eps], device: device) - self.fused = ProcessInfo.processInfo.environment["FFAI_GDN_FUSED_PREP"] != nil + // Fused-prep is DEFAULT ON for Qwen3.6-A3B. The legacy host-loop + // path lives behind `FFAI_GDN_NO_FUSED_PREP=1` for precision + // comparison. + self.fused = + ProcessInfo.processInfo.environment["FFAI_GDN_NO_FUSED_PREP"] == nil // Per-decode-token scratch — pre-allocated once at init so the // fused GDN path doesn't pay 6 × MTLBuffer allocations per call. - // See the comments next to the field declarations for the - // cross-token safety argument (Metal hazard tracking + the - // engine's per-token cmd wait). self.convOutScratch = Tensor.empty(shape: [convDim], dtype: dtype, device: device) self.convActF32Scratch = Tensor.empty(shape: [convDim], dtype: .f32, device: device) self.aRawF32Scratch = Tensor.empty(shape: [numValueHeads], dtype: .f32, device: device) @@ -1148,6 +1289,57 @@ public final class Qwen35GDNMixer: Module { shape: [numValueHeads, valueHeadDim], dtype: .f32, device: device) self.yGatedScratch = Tensor.empty(shape: [valueDim], dtype: dtype, device: device) + + // Set up the batched 4-projection input fast path when all 4 + // inProj are QuantizedLinear int4 with the same groupSize, and + // pin scratches in the device residency set so Metal skips + // per-encoder residency revalidation. + if let qQL = inProjQKV.inner as? QuantizedLinear, + let zQL = inProjZ.inner as? QuantizedLinear, + let bQL = inProjB.inner as? QuantizedLinear, + let aQL = inProjA.inner as? QuantizedLinear, + qQL.bits == 4 && zQL.bits == 4 && bQL.bits == 4 && aQL.bits == 4, + qQL.groupSize == zQL.groupSize, + qQL.groupSize == bQL.groupSize, + qQL.groupSize == aQL.groupSize + { + self.batchedInputProj = ( + qW: qQL.weight, qS: qQL.scales, qB: qQL.biases, + zW: zQL.weight, zS: zQL.scales, zB: zQL.biases, + bW: bQL.weight, bS: bQL.scales, bB: bQL.biases, + aW: aQL.weight, aS: aQL.scales, aB: aQL.biases, + groupSize: qQL.groupSize + ) + // Single-dispatch fast path eligibility. + let inDim = qQL.weight.shape[1] * 8 + if inDim % 512 == 0 + && convDim % 8 == 0 && valueDim % 8 == 0 + && numValueHeads % 8 == 0 + && qQL.groupSize == 64 + && ProcessInfo.processInfo.environment["FFAI_NO_FUSED_GDN_4"] == nil + { + self.fused4Eligible = true + } + self.qkvScratch = Tensor.empty( + shape: [convDim], dtype: dtype, device: device) + self.zScratch = Tensor.empty( + shape: [valueDim], dtype: dtype, device: device) + self.bRawScratch = Tensor.empty( + shape: [numValueHeads], dtype: dtype, device: device) + self.aRawScratch = Tensor.empty( + shape: [numValueHeads], dtype: dtype, device: device) + // Pin scratch buffers + fused-path constants so Metal + // skips per-encoder residency revalidation. + device.markWeightsResident([ + qkvScratch!.buffer, zScratch!.buffer, + bRawScratch!.buffer, aRawScratch!.buffer, + convOutScratch.buffer, convActF32Scratch.buffer, + aRawF32Scratch.buffer, bRawF32Scratch.buffer, + yF32Scratch.buffer, yGatedScratch.buffer, + qNormWeightF32.buffer, kNormWeightF32.buffer, + aLogTF32.buffer, dtBiasTF32.buffer, epsBufFused.buffer, + ]) + } } public func parameters() -> [(String, Tensor)] { @@ -1182,16 +1374,62 @@ public final class Qwen35GDNMixer: Module { cmd: MTLCommandBuffer, device: Device ) -> Tensor { // ── GPU phase 1: projections + conv + SiLU ──────────────────── - let qkv = inProjQKV(xNorm, on: cmd) // [conv_dim] - let z = inProjZ(xNorm, on: cmd) // [value_dim] - let bRaw = inProjB(xNorm, on: cmd) // [num_value_heads] - let aRaw = inProjA(xNorm, on: cmd) // [num_value_heads] + // 4-projection shared-encoder fast path when all inProj are int4 + // QuantizedLinear with the same groupSize. Saves 3 encoder + // begin/end pairs per GDN layer × 30 ≈ 1.5 ms/token. Upgrades + // to a single-dispatch fast path via `batched4QgemvInt4Fast` + // when constraints match (out_dim%8, in_dim%512, groupSize=64). + let qkv: Tensor + let z: Tensor + let bRaw: Tensor + let aRaw: Tensor + if let bp = batchedInputProj { + qkv = qkvScratch! + z = zScratch! + bRaw = bRawScratch! + aRaw = aRawScratch! + if fused4Eligible { + Ops.batched4QgemvInt4Fast( + input: xNorm, + wA: bp.qW, scalesA: bp.qS, biasesA: bp.qB, outA: qkv, + wB: bp.zW, scalesB: bp.zS, biasesB: bp.zB, outB: z, + wC: bp.bW, scalesC: bp.bS, biasesC: bp.bB, outC: bRaw, + wD: bp.aW, scalesD: bp.aS, biasesD: bp.aB, outD: aRaw, + groupSize: bp.groupSize, on: cmd) + } else { + Ops.dequantGemvInt4Four( + input: xNorm, + w0: bp.qW, s0: bp.qS, b0: bp.qB, out0: qkv, + w1: bp.zW, s1: bp.zS, b1: bp.zB, out1: z, + w2: bp.bW, s2: bp.bS, b2: bp.bB, out2: bRaw, + w3: bp.aW, s3: bp.aS, b3: bp.aB, out3: aRaw, + groupSize: bp.groupSize, on: cmd) + } + } else { + qkv = inProjQKV(xNorm, on: cmd) // [conv_dim] + z = inProjZ(xNorm, on: cmd) // [value_dim] + bRaw = inProjB(xNorm, on: cmd) // [num_value_heads] + aRaw = inProjA(xNorm, on: cmd) // [num_value_heads] + } Ops.conv1dCausalStep( x: qkv, w: convW, b: convB, state: cache.conv.state, into: convOutScratch, nChannels: convDim, kernelSize: convKernel, on: cmd) - let convAct = Ops.silu(convOutScratch, on: cmd) // [conv_dim] + // In the fused path the convOut silu + cast to f32 fuse with + // the aRaw / bRaw casts via `siluCastF32PlusCastF32Two` later + // on the same shared encoder. Outside the fused path the + // legacy host phase still needs the bf16 silu output, so + // compute it eagerly there. + let convAct: Tensor + if fused { + // Fused path: skip the silu here; the combined + // siluCast+cast2 dispatch below produces the f32 input + // directly. + convAct = convOutScratch + } else { + convAct = Ops.silu(convOutScratch, on: cmd) // [conv_dim] + } // ── Fused GDN prep + recurrence path (opt-in) ───────────────── // @@ -1218,9 +1456,15 @@ public final class Qwen35GDNMixer: Module { // fused step runs the recurrence in fp32 against the // existing fp32 state slots, matching the canonical // precision of the legacy path. - Ops.castToF32(convAct, into: convActF32Scratch, on: cmd) - Ops.castToF32(aRaw, into: aRawF32Scratch, on: cmd) - Ops.castToF32(bRaw, into: bRawF32Scratch, on: cmd) + // Collapse silu+cast (convOut → f32) AND the two plain + // casts (aRaw, bRaw → f32) onto ONE shared encoder. Saves + // 1 encoder begin/end per GDN layer × 30 ≈ 30 pairs per + // decode token. + Ops.siluCastF32PlusCastF32Two( + siluIn: convOutScratch, into: convActF32Scratch, + aRaw, into: aRawF32Scratch, + bRaw, into: bRawF32Scratch, + on: cmd) Ops.gatedDeltaPrepStep( convOut: convActF32Scratch, @@ -1406,6 +1650,15 @@ public final class Qwen35GDNMixer: Module { "Qwen35GDNMixer.forwardMany: xNormFlat size \(xNormFlat.elementCount) ≠ T·hidden = \(t * hidden)" ) + // Chunked prep+recurrence is the default route. Replaces the + // per-token T-loop body with one `Ops.gatedDeltaPrepChunk` + // dispatch — state register-resident across the full T-sweep. + // Opt out via FFAI_GDN_NO_PREP_CHUNK=1 for A/B benching. + if ProcessInfo.processInfo.environment["FFAI_GDN_NO_PREP_CHUNK"] == nil { + return forwardManyChunked( + xNormFlat, t: t, cache: cache, cmd: cmd, device: device) + } + let dt = xNormFlat.dtype let dtBytes = dt.byteSize @@ -1415,8 +1668,20 @@ public final class Qwen35GDNMixer: Module { let bRawAll = inProjB.callMany(xNormFlat, t: t, on: cmd, device: device) let aRawAll = inProjA.callMany(xNormFlat, t: t, on: cmd, device: device) + // a_raw / b_raw → f32 once via shared-encoder `castToF32Two`, + // then per-row slices alias into the pre-cast f32 storage. + let aRawF32All = Tensor.empty( + shape: [t * numValueHeads], dtype: .f32, device: device) + let bRawF32All = Tensor.empty( + shape: [t * numValueHeads], dtype: .f32, device: device) + Ops.castToF32Two( + aRawAll, into: aRawF32All, + bRawAll, into: bRawF32All, + on: cmd) + // ── Per-token recurrence — scratches reused, T-loop on `cmd`. ──── let yGatedAll = Tensor.empty(shape: [t * valueDim], dtype: dt, device: device) + let f32Bytes = DType.f32.byteSize for r in 0 ..< t { let qkvRow = Tensor( buffer: qkvAll.buffer, @@ -1426,38 +1691,30 @@ public final class Qwen35GDNMixer: Module { buffer: zAll.buffer, offset: zAll.offset + r * valueDim * dtBytes, shape: [valueDim], dtype: dt) - let bRawRow = Tensor( - buffer: bRawAll.buffer, - offset: bRawAll.offset + r * numValueHeads * dtBytes, - shape: [numValueHeads], dtype: dt) - let aRawRow = Tensor( - buffer: aRawAll.buffer, - offset: aRawAll.offset + r * numValueHeads * dtBytes, - shape: [numValueHeads], dtype: dt) + let aRawF32Row = Tensor( + buffer: aRawF32All.buffer, + offset: aRawF32All.offset + r * numValueHeads * f32Bytes, + shape: [numValueHeads], dtype: .f32) + let bRawF32Row = Tensor( + buffer: bRawF32All.buffer, + offset: bRawF32All.offset + r * numValueHeads * f32Bytes, + shape: [numValueHeads], dtype: .f32) Ops.conv1dCausalStep( x: qkvRow, w: convW, b: convB, state: cache.conv.state, into: convOutScratch, nChannels: convDim, kernelSize: convKernel, on: cmd) - // Fused silu + cast-to-f32. Replaces the prior two-dispatch - // chain `silu(convOutScratch) → castToF32(...)` with one - // dispatch reading bf16/f16 conv output, applying silu at - // fp32 precision, writing fp32 directly into the prep - // scratch. Saves T·30 dispatches per Qwen3.6-A3B prefill. if convOutScratch.dtype == .f32 { - // f32 conv → silu in place, no cast needed. _ = Ops.silu(convOutScratch, on: cmd, into: convOutScratch) Ops.castToF32(convOutScratch, into: convActF32Scratch, on: cmd) } else { Ops.siluCastToF32(convOutScratch, into: convActF32Scratch, on: cmd) } - Ops.castToF32(aRawRow, into: aRawF32Scratch, on: cmd) - Ops.castToF32(bRawRow, into: bRawF32Scratch, on: cmd) Ops.gatedDeltaPrepStep( convOut: convActF32Scratch, aLog: aLogTF32, dtBias: dtBiasTF32, - aRaw: aRawF32Scratch, bRaw: bRawF32Scratch, + aRaw: aRawF32Row, bRaw: bRawF32Row, qNormWeight: qNormWeightF32, kNormWeight: kNormWeightF32, stateIn: cache.gdn.current, stateOut: cache.gdn.next, y: yF32Scratch, @@ -1483,6 +1740,135 @@ public final class Qwen35GDNMixer: Module { let yGatedRows = yGatedAll.reshaped(to: [t, valueDim]) return outProj.callMany(yGatedRows, t: t, on: cmd, device: device) } + + /// Chunked GDN forward — replaces the per-token T-loop body with one + /// `Ops.gatedDeltaPrepChunk` dispatch. The recurrence state stays in + /// per-lane registers across all T tokens; the kernel loads state + /// once, sweeps T tokens, stores state once. State traffic per layer + /// drops from `T × (load + store)` to `1 × (load + store)`. + /// + /// Conv1d still loops T times unless the batched + /// `conv1dCausalStepSiluCastMany` kernel applies (convKernel == 4), + /// in which case all T conv1d_steps fuse with silu+cast into ONE + /// dispatch and conv state stays in per-channel registers. + /// + /// Mixer norm runs once over T·Hv rows via `gatedMixerNormMany`. + /// Output projection fans through a batched qmm. Gated by the + /// `forwardMany` dispatcher; opt out via `FFAI_GDN_NO_PREP_CHUNK=1`. + func forwardManyChunked( + _ xNormFlat: Tensor, t: Int, + cache: Qwen35GDNLayerCache, + cmd: MTLCommandBuffer, device: Device + ) -> Tensor { + let dt = xNormFlat.dtype + let dtBytes = dt.byteSize + let f32Bytes = DType.f32.byteSize + let _ = f32Bytes // silence unused warning when conv-many path used + + // ── Five T-batched projections (same as forwardMany) ───────────── + let qkvAll = inProjQKV.callMany(xNormFlat, t: t, on: cmd, device: device) + let zAll = inProjZ.callMany(xNormFlat, t: t, on: cmd, device: device) + let bRawAll = inProjB.callMany(xNormFlat, t: t, on: cmd, device: device) + let aRawAll = inProjA.callMany(xNormFlat, t: t, on: cmd, device: device) + + // a_raw / b_raw → f32 once on ONE shared encoder via + // `castToF32Two`. Saves one encoder begin/end pair per GDN + // layer × 30 = 30 pairs/prefill. + let aRawF32All = Tensor.empty( + shape: [t * numValueHeads], dtype: .f32, device: device) + let bRawF32All = Tensor.empty( + shape: [t * numValueHeads], dtype: .f32, device: device) + Ops.castToF32Two( + aRawAll, into: aRawF32All, + bRawAll, into: bRawF32All, + on: cmd) + + // ── Conv1d + silu+cast → `[T, convDim]` f32 staging buffer ────── + // When conv_kernel = 4, the batched + // `conv1dCausalStepSiluCastMany` kernel sweeps all T tokens in + // ONE dispatch with the conv state held in per-channel registers + // across the sweep — state read once at start, written once at + // end. Saves T-1 dispatches per layer × 30 ≈ 15K dispatches at + // T=512, and state DRAM traffic drops to 1·(load+store). Gate + // via FFAI_NO_CONV1D_MANY=1. + let convOutAllF32 = Tensor.empty( + shape: [t * convDim], dtype: .f32, device: device) + if convKernel == 4 + && ProcessInfo.processInfo.environment["FFAI_NO_CONV1D_MANY"] == nil + { + Ops.conv1dCausalStepSiluCastMany( + src: qkvAll.reshaped(to: [t, convDim]), + w: convW, b: convB, + stateIn: cache.conv.state, + outF32: convOutAllF32, + stateOut: cache.conv.state, + t: t, convDim: convDim, convKernel: convKernel, + on: cmd) + } else { + for r in 0 ..< t { + let qkvRow = Tensor( + buffer: qkvAll.buffer, + offset: qkvAll.offset + r * convDim * dtBytes, + shape: [convDim], dtype: dt) + Ops.conv1dCausalStep( + x: qkvRow, w: convW, b: convB, + state: cache.conv.state, into: convOutScratch, + nChannels: convDim, kernelSize: convKernel, on: cmd) + let convOutRowF32 = Tensor( + buffer: convOutAllF32.buffer, + offset: convOutAllF32.offset + r * convDim * f32Bytes, + shape: [convDim], dtype: .f32) + if convOutScratch.dtype == .f32 { + _ = Ops.silu(convOutScratch, on: cmd, into: convOutScratch) + Ops.castToF32(convOutScratch, into: convOutRowF32, on: cmd) + } else { + Ops.siluCastToF32( + convOutScratch, into: convOutRowF32, on: cmd) + } + } + } + + // ── ONE chunked prep+recurrence dispatch over all T tokens ────── + let tLenBuf = device.makeBuffer(length: 4) + var tLenU32 = UInt32(t) + memcpy(tLenBuf.contents(), &tLenU32, 4) + let tLenScalar = Tensor( + buffer: tLenBuf, offset: 0, shape: [1], dtype: .u32) + + let yF32All = Tensor.empty( + shape: [t * numValueHeads * valueHeadDim], + dtype: .f32, device: device) + + Ops.gatedDeltaPrepChunk( + convOut: convOutAllF32, + aLog: aLogTF32, dtBias: dtBiasTF32, + aRaw: aRawF32All, bRaw: bRawF32All, + qNormWeight: qNormWeightF32, kNormWeight: kNormWeightF32, + stateIn: cache.gdn.current, stateOut: cache.gdn.next, + y: yF32All, + tLen: tLenScalar, + batchSize: 1, dk: keyHeadDim, dv: valueHeadDim, + hv: numValueHeads, hk: numKeyHeads, + on: cmd) + cache.gdn.swap() + for _ in 0 ..< t { cache.advance() } + + // ── Mixer norm T-batched (one dispatch across T·Hv rows) ──────── + // `gatedMixerNormMany` collapses T per-row dispatches into one + // shared-encoder dispatch. + let yGatedAll = Tensor.empty( + shape: [t * valueDim], dtype: dt, device: device) + Ops.gatedMixerNormMany( + y: yF32All, z: zAll, weight: mixerNorm.weight, + epsBuf: epsBufFused, + into: yGatedAll, + t: t, numValueHeads: numValueHeads, valueHeadDim: valueHeadDim, + on: cmd) + + // ── Batched output projection ─────────────────────────────────── + let yGatedRows = yGatedAll.reshaped(to: [t, valueDim]) + return outProj.callMany(yGatedRows, t: t, on: cmd, device: device) + } } // ─── Qwen35AttentionMixer — gated multi-head attention ─────────────── @@ -1491,6 +1877,13 @@ public final class Qwen35GDNMixer: Module { // gate); the attention output is multiplied by `sigmoid(gate)` before // `o_proj`. `q_norm` / `k_norm` are per-head RMSNorm. RoPE is partial // (`partial_rotary_factor`): only the first `rotaryDim` dims rotate. +// +// Threading contract: instances are NOT safe for concurrent calls. +// `forward` reuses ~12 per-instance scratch tensors (`qOutScratch`, +// `qNormedScratch`, `attnOutScratch`, `partialOScratch`, …) across +// calls without internal synchronisation. `forwardMany` allocates +// fresh per-call tensors and does not touch the single-token +// scratches. Callers must serialise inference on a given instance. public final class Qwen35AttentionMixer: Module { let qProj, kProj, vProj, oProj: AnyLinear @@ -1499,12 +1892,64 @@ public final class Qwen35AttentionMixer: Module { let ropeTheta: Float let attnOutputGate: Bool let scale: Float + /// Cached row-index tensors for sliceHeadHalves35. The indices are + /// constant per `nHeads` (even rows = queries, odd rows = gates) — + /// previously rebuilt + memcpy'd into a fresh MTLBuffer per attn + /// layer per decode token. + let sliceFirstIdx: Tensor? // [nHeads] u32 — 0, 2, 4, ... + let sliceSecondIdx: Tensor? // [nHeads] u32 — 1, 3, 5, ... + /// Batched QKV-projection fast-path detection. When all 3 + /// projections (q, k, v) are QuantizedLinear int4 with the same + /// groupSize, the mixer routes through `Ops.dequantGemvInt4Three` + /// — 3 qmms on one shared encoder. + private var batchedQKV: + ( + qW: Tensor, qS: Tensor, qB: Tensor, + kW: Tensor, kS: Tensor, kB: Tensor, + vW: Tensor, vS: Tensor, vB: Tensor, + groupSize: Int + )? + private var qOutScratch: Tensor? + private var kOutScratch: Tensor? + private var vOutScratch: Tensor? + /// Fused-QKV single-dispatch fast path. When set, `forward` uses + /// `Ops.batchedQkvQgemvInt4Fast` (1 dispatch) instead of + /// `dequantGemvInt4Three` (3 dispatches on shared encoder). + /// Backing scratch is one contiguous `[qDim+kDim+vDim]` buffer. + private var fusedQKVEligible: Bool = false + private var qkvFusedScratch: Tensor? + /// Cached q_norm + k_norm outputs (was fresh per call). + private var qNormedScratch: Tensor? + private var kNormedScratch: Tensor? + /// Cached sliceHeadHalves35 gather outputs. + private var queriesScratch: Tensor? + private var gateSliceScratch: Tensor? + /// Cached sdpaDecode output + sigmoidMul output. + private var attnOutScratch: Tensor? + private var attnFlatGatedScratch: Tensor? + /// Cached Flash-Decoding 2-pass partial buffers. Allocated lazily + /// on first 2-pass dispatch (when kv.length ≥ threshold). Sized for + /// (nHeads, blocks=32, headDim) — ~192KB per attn layer on + /// Qwen3.6 bf16. + private var partialOScratch: Tensor? + private var partialMScratch: Tensor? + private var partialLScratch: Tensor? + private let sdpa2PassBlocks: Int = 32 + private let sdpa2PassThreshold: Int = { + if let raw = ProcessInfo.processInfo.environment["FFAI_SDPA_2PASS_THRESHOLD"], + let n = Int(raw) + { + return n + } + return 1024 // default: only kick in at long KV + }() init( qProj: AnyLinear, kProj: AnyLinear, vProj: AnyLinear, oProj: AnyLinear, qNorm: RMSNorm, kNorm: RMSNorm, nHeads: Int, nKVHeads: Int, headDim: Int, rotaryDim: Int, - ropeTheta: Float, attnOutputGate: Bool + ropeTheta: Float, attnOutputGate: Bool, + device: Device = .shared ) { self.qProj = qProj self.kProj = kProj @@ -1519,6 +1964,61 @@ public final class Qwen35AttentionMixer: Module { self.ropeTheta = ropeTheta self.attnOutputGate = attnOutputGate self.scale = 1.0 / Float(Double(headDim).squareRoot()) + if attnOutputGate { + // Pre-build the row-index tensors for the gate-split + // gather. Even rows = queries, odd rows = gates. + var firstRows = [UInt32](repeating: 0, count: nHeads) + var secondRows = [UInt32](repeating: 0, count: nHeads) + for h in 0 ..< nHeads { + firstRows[h] = UInt32(2 * h) + secondRows[h] = UInt32(2 * h + 1) + } + let firstBuf = device.makeBuffer(length: nHeads * 4) + let secondBuf = device.makeBuffer(length: nHeads * 4) + firstRows.withUnsafeBytes { + _ = memcpy(firstBuf.contents(), $0.baseAddress!, nHeads * 4) + } + secondRows.withUnsafeBytes { + _ = memcpy(secondBuf.contents(), $0.baseAddress!, nHeads * 4) + } + self.sliceFirstIdx = Tensor( + buffer: firstBuf, offset: 0, shape: [nHeads], dtype: .u32) + self.sliceSecondIdx = Tensor( + buffer: secondBuf, offset: 0, shape: [nHeads], dtype: .u32) + // Pin idx scratches to the residency set. + device.markWeightsResident([firstBuf, secondBuf]) + } else { + self.sliceFirstIdx = nil + self.sliceSecondIdx = nil + } + // Detect batched-qkv fast path. + if let qQL = qProj.inner as? QuantizedLinear, + let kQL = kProj.inner as? QuantizedLinear, + let vQL = vProj.inner as? QuantizedLinear, + qQL.bits == 4 && kQL.bits == 4 && vQL.bits == 4, + qQL.groupSize == kQL.groupSize, qQL.groupSize == vQL.groupSize + { + self.batchedQKV = ( + qW: qQL.weight, qS: qQL.scales, qB: qQL.biases, + kW: kQL.weight, kS: kQL.scales, kB: kQL.biases, + vW: vQL.weight, vS: vQL.scales, vB: vQL.biases, + groupSize: qQL.groupSize + ) + // Fused single-dispatch QKV gemv. Constraints: in_dim + // multiple of 512, each out_dim multiple of 8, + // group_size=64. + let inDim = qQL.weight.shape[1] * 8 + let qDim = qQL.weight.shape[0] + let kDim = kQL.weight.shape[0] + let vDim = vQL.weight.shape[0] + if inDim % 512 == 0 + && qDim % 8 == 0 && kDim % 8 == 0 && vDim % 8 == 0 + && qQL.groupSize == 64 + && ProcessInfo.processInfo.environment["FFAI_NO_FUSED_QKV"] == nil + { + self.fusedQKVEligible = true + } + } } public func parameters() -> [(String, Tensor)] { @@ -1542,42 +2042,127 @@ public final class Qwen35AttentionMixer: Module { // q_proj projects 2× heads when attn_output_gate is set: the // first `nHeads · headDim` elements are the queries, the second // half is the per-head sigmoid gate. - let qOut = qProj(xNorm, on: cmd) + // Batch q+k+v projections into one encoder when all 3 are + // QuantizedLinear int4; when the fast-path constraints match, + // fuse all three into a single dispatch. + let qOut: Tensor + let kPre: Tensor + let vPre: Tensor + if let bq = batchedQKV { + let qDim = bq.qW.shape[0] + let kDim = bq.kW.shape[0] + let vDim = bq.vW.shape[0] + if fusedQKVEligible { + // ONE dispatch into a [qDim+kDim+vDim] scratch. + // qOut / kPre / vPre are offset views into the same buffer. + let total = qDim + kDim + vDim + if qkvFusedScratch == nil + || qkvFusedScratch!.elementCount != total + || qkvFusedScratch!.dtype != xNorm.dtype + { + qkvFusedScratch = Tensor.empty( + shape: [total], dtype: xNorm.dtype, device: device) + qOutScratch = qkvFusedScratch!.slicedRows(start: 0, count: qDim) + kOutScratch = qkvFusedScratch!.slicedRows(start: qDim, count: kDim) + vOutScratch = qkvFusedScratch!.slicedRows( + start: qDim + kDim, count: vDim) + } + qOut = qOutScratch! + kPre = kOutScratch! + vPre = vOutScratch! + Ops.batchedQkvQgemvInt4Fast( + x: xNorm, + wQ: bq.qW, scalesQ: bq.qS, biasesQ: bq.qB, + wK: bq.kW, scalesK: bq.kS, biasesK: bq.kB, + wV: bq.vW, scalesV: bq.vS, biasesV: bq.vB, + outQ: qDim, outK: kDim, outV: vDim, + on: cmd, into: qkvFusedScratch!) + } else { + // Legacy 3-dispatch shared-encoder path. + if qOutScratch == nil || qOutScratch!.elementCount != qDim { + qOutScratch = Tensor.empty( + shape: [qDim], dtype: xNorm.dtype, device: device) + } + if kOutScratch == nil || kOutScratch!.elementCount != kDim { + kOutScratch = Tensor.empty( + shape: [kDim], dtype: xNorm.dtype, device: device) + } + if vOutScratch == nil || vOutScratch!.elementCount != vDim { + vOutScratch = Tensor.empty( + shape: [vDim], dtype: xNorm.dtype, device: device) + } + qOut = qOutScratch! + kPre = kOutScratch! + vPre = vOutScratch! + Ops.dequantGemvInt4Three( + input: xNorm, + w0: bq.qW, s0: bq.qS, b0: bq.qB, out0: qOut, + w1: bq.kW, s1: bq.kS, b1: bq.kB, out1: kPre, + w2: bq.vW, s2: bq.vS, b2: bq.vB, out2: vPre, + groupSize: bq.groupSize, on: cmd) + } + } else { + qOut = qProj(xNorm, on: cmd) + kPre = kProj(xNorm, on: cmd) + vPre = vProj(xNorm, on: cmd) + } let queries: Tensor let gate: Tensor? if attnOutputGate { // Layout is [nHeads, 2 · headDim] — per head the first // `headDim` is the query, the next `headDim` the gate. + // Use cached idx tensors. let q2 = qOut.reshaped(to: [nHeads, 2 * headDim]) - queries = sliceHeadHalves35( - q2, nHeads: nHeads, headDim: headDim, takeFirst: true, - on: cmd, device: device) - gate = sliceHeadHalves35( - q2, nHeads: nHeads, headDim: headDim, takeFirst: false, - on: cmd, device: device) + let table = q2.reshaped(to: [nHeads * 2, headDim]) + if queriesScratch == nil + || queriesScratch!.elementCount != nHeads * headDim + || queriesScratch!.dtype != qOut.dtype + { + queriesScratch = Tensor.empty( + shape: [nHeads, headDim], dtype: qOut.dtype, device: device) + gateSliceScratch = Tensor.empty( + shape: [nHeads, headDim], dtype: qOut.dtype, device: device) + } + // Batched two-gather on shared encoder, saves 1 encoder + // begin/end pair per attn layer × 10 = ~170 µs/token. + Ops.gatherTwo( + table: table, + ids1: sliceFirstIdx!, into: queriesScratch!, + ids2: sliceSecondIdx!, into: gateSliceScratch!, + on: cmd) + queries = queriesScratch!.reshaped(to: [nHeads * headDim]) + gate = gateSliceScratch!.reshaped(to: [nHeads * headDim]) } else { queries = qOut gate = nil } - let k = kProj(xNorm, on: cmd) - let v = vProj(xNorm, on: cmd) + let k = kPre + let v = vPre // Per-head q_norm / k_norm (weighted RMSNorm over headDim). - let qNormed = Ops.rmsNormRows( - queries, weight: qNorm.weight, eps: qNorm.eps, - nRows: nHeads, rowSize: headDim, on: cmd) - let kNormed = Ops.rmsNormRows( - k, weight: kNorm.weight, eps: kNorm.eps, - nRows: nKVHeads, rowSize: headDim, on: cmd) + // Shared encoder + cached scratches. + if qNormedScratch == nil + || qNormedScratch!.elementCount != queries.elementCount + || qNormedScratch!.dtype != queries.dtype + { + qNormedScratch = Tensor.empty( + shape: queries.shape, dtype: queries.dtype, device: device) + kNormedScratch = Tensor.empty( + shape: k.shape, dtype: k.dtype, device: device) + } + let qNormed = qNormedScratch! + let kNormed = kNormedScratch! + Ops.rmsNormRowsTwo( + queries, weight: qNorm.weight, eps1: qNorm.eps, + nRows1: nHeads, rowSize1: headDim, into: qNormed, + k, weight: kNorm.weight, eps2: kNorm.eps, + nRows2: nKVHeads, rowSize2: headDim, into: kNormed, + on: cmd) // Partial RoPE — rotate only the first `rotaryDim` dims of each - // head, in place. - Ops.ropePartial( - qNormed, position: position, - headDim: headDim, rotaryDim: rotaryDim, - thetaBase: ropeTheta, on: cmd) - Ops.ropePartial( - kNormed, position: position, + // head, in place. q + k on shared encoder. + Ops.ropePartialTwo( + qNormed, kNormed, position: position, headDim: headDim, rotaryDim: rotaryDim, thetaBase: ropeTheta, on: cmd) @@ -1587,16 +2172,60 @@ public final class Qwen35AttentionMixer: Module { vFlat: v.reshaped(to: [nKVHeads, headDim]), on: cmd) let (cacheK, cacheV) = kv.prepareForAttention(on: cmd) - let attnOut = Ops.sdpaDecode( - q: qNormed.reshaped(to: [nHeads, headDim]), k: cacheK, v: cacheV, - nQHeads: nHeads, nKVHeads: nKVHeads, headDim: headDim, - nKV: kv.length, kvStride: kv.maxSeq, - scale: scale, on: cmd) + // Cached sdpaDecode output + sigmoidMul output scratches. + if attnOutScratch == nil + || attnOutScratch!.elementCount != nHeads * headDim + || attnOutScratch!.dtype != qNormed.dtype + { + attnOutScratch = Tensor.empty( + shape: [nHeads, headDim], dtype: qNormed.dtype, device: device) + attnFlatGatedScratch = Tensor.empty( + shape: [nHeads * headDim], dtype: qNormed.dtype, device: device) + } + // Flash-Decoding 2-pass at long KV. Splits KV across `blocks` + // threadgroups in pass1, merges in pass2. Wins over single-pass + // when nKV ≥ ~1024 (default threshold); env var + // `FFAI_SDPA_2PASS_THRESHOLD` overrides. + let attnOut: Tensor + if kv.length >= sdpa2PassThreshold { + if partialOScratch == nil + || partialOScratch!.elementCount + != nHeads * sdpa2PassBlocks * headDim + || partialOScratch!.dtype != qNormed.dtype + { + partialOScratch = Tensor.empty( + shape: [nHeads, sdpa2PassBlocks, headDim], + dtype: qNormed.dtype, device: device) + partialMScratch = Tensor.empty( + shape: [nHeads, sdpa2PassBlocks], dtype: .f32, device: device) + partialLScratch = Tensor.empty( + shape: [nHeads, sdpa2PassBlocks], dtype: .f32, device: device) + } + Ops.sdpaDecode2Pass( + q: qNormed.reshaped(to: [nHeads, headDim]), k: cacheK, v: cacheV, + nQHeads: nHeads, nKVHeads: nKVHeads, headDim: headDim, + nKV: kv.length, kvStride: kv.maxSeq, blocks: sdpa2PassBlocks, + scale: scale, + partialO: partialOScratch!, + partialM: partialMScratch!, + partialL: partialLScratch!, + into: attnOutScratch!, on: cmd) + attnOut = attnOutScratch! + } else { + attnOut = Ops.sdpaDecode( + q: qNormed.reshaped(to: [nHeads, headDim]), k: cacheK, v: cacheV, + nQHeads: nHeads, nKVHeads: nKVHeads, headDim: headDim, + nKV: kv.length, kvStride: kv.maxSeq, + scale: scale, on: cmd, into: attnOutScratch) + } // Gated output: attnOut * sigmoid(gate), then o_proj. + // Fused sigmoidMul via mt_sigmoid_mul into cached scratch — + // saves 1 encoder per attn layer × 10 ≈ 170 µs/token. var attnFlat = attnOut.reshaped(to: [nHeads * headDim]) if let gate { - attnFlat = Ops.mul(attnFlat, Ops.sigmoid(gate, on: cmd), on: cmd) + attnFlat = Ops.sigmoidMul( + attnFlat, gate, on: cmd, into: attnFlatGatedScratch) } return oProj(attnFlat, on: cmd) } @@ -1632,101 +2261,146 @@ public final class Qwen35AttentionMixer: Module { let qDim = nHeads * headDim let kvDim = nKVHeads * headDim - // ── Projections — one gemm/qmm each, T-batched ─────────────────── - let qOut = qProj.callMany(xNormFlat, t: t, on: cmd, device: device) - let kOut = kProj.callMany(xNormFlat, t: t, on: cmd, device: device) - let vOut = vProj.callMany(xNormFlat, t: t, on: cmd, device: device) + // ── Projections — fused QKV qmm when constraints match ────────── + // `Ops.batchedQkvQmmFast` produces all three Q/K/V projections + // in ONE dispatch when the three QuantizedLinear projections + // share `bits=4`, `groupSize=64`, and `in_dim%512==0` / + // `out_*%8==0`. Replaces 3 `callMany` qmm encoders with 1 fused + // encoder. Predicate is the same `fusedQKVEligible` flag the + // single-token path uses. + let qOut: Tensor + let kOut: Tensor + let vOut: Tensor + if fusedQKVEligible, let bq = batchedQKV { + let qDimW = bq.qW.shape[0] + let kDimW = bq.kW.shape[0] + let vDimW = bq.vW.shape[0] + let inDim = bq.qW.shape[1] * 8 + let qBuf = Tensor.empty(shape: [t, qDimW], dtype: dt, device: device) + let kBuf = Tensor.empty(shape: [t, kDimW], dtype: dt, device: device) + let vBuf = Tensor.empty(shape: [t, vDimW], dtype: dt, device: device) + Ops.batchedQkvQmmFast( + x: xNormFlat.reshaped(to: [t, inDim]), + wQ: bq.qW, scalesQ: bq.qS, biasesQ: bq.qB, + wK: bq.kW, scalesK: bq.kS, biasesK: bq.kB, + wV: bq.vW, scalesV: bq.vS, biasesV: bq.vB, + m: t, outQ: qDimW, outK: kDimW, outV: vDimW, + on: cmd, + qBuf: qBuf, kBuf: kBuf, vBuf: vBuf) + qOut = qBuf + kOut = kBuf + vOut = vBuf + } else { + qOut = qProj.callMany(xNormFlat, t: t, on: cmd, device: device) + kOut = kProj.callMany(xNormFlat, t: t, on: cmd, device: device) + vOut = vProj.callMany(xNormFlat, t: t, on: cmd, device: device) + } - // ── Gate split — one gather (vs T) when attnOutputGate ─────────── + // ── Gate split — one shared-encoder gather (vs T) when + // attnOutputGate. Both halves run on a shared encoder via + // sliceHeadHalvesManyBoth35. let queriesFlat: Tensor // `[T, nHeads, headDim]` flat let gateFlat: Tensor? // `[T, nHeads, headDim]` flat if attnOutputGate { let q2T = qOut.reshaped(to: [t, nHeads, 2 * headDim]) - queriesFlat = sliceHeadHalvesMany35( - q2T, t: t, nHeads: nHeads, headDim: headDim, takeFirst: true, - on: cmd, device: device) - gateFlat = sliceHeadHalvesMany35( - q2T, t: t, nHeads: nHeads, headDim: headDim, takeFirst: false, + let (q, g) = sliceHeadHalvesManyBoth35( + q2T, t: t, nHeads: nHeads, headDim: headDim, on: cmd, device: device) + queriesFlat = q + gateFlat = g } else { queriesFlat = qOut gateFlat = nil } - // ── Q/K norm — one rmsNormRows over (T*nHeads) and (T*nKVHeads) - // rows. rmsNormRows is already row-wise, so the T-batched form is - // a count change, not a kernel change. - let qNormed = Ops.rmsNormRows( - queriesFlat, weight: qNorm.weight, eps: qNorm.eps, - nRows: t * nHeads, rowSize: headDim, on: cmd) - let kNormed = Ops.rmsNormRows( - kOut, weight: kNorm.weight, eps: kNorm.eps, - nRows: t * nKVHeads, rowSize: headDim, on: cmd) - - // ── RoPE per token + KV append — T dispatches each, all on `cmd`. - // The single-token RoPE kernel grids `[nHeads, halfRotary]` — - // small enough that T launches at T≤256 cost less than one big - // qmm. KV append uses appendRangeOnGPU for the single length-lock - // path; each step is one Ops.kvCacheUpdate dispatch. - var kRows: [Tensor] = [] - kRows.reserveCapacity(t) - var vRows: [Tensor] = [] - vRows.reserveCapacity(t) - for r in 0 ..< t { - let qRow = Tensor( - buffer: qNormed.buffer, - offset: qNormed.offset + r * qDim * dtBytes, - shape: [qDim], dtype: dt) - let kRow = Tensor( - buffer: kNormed.buffer, - offset: kNormed.offset + r * kvDim * dtBytes, - shape: [kvDim], dtype: dt) - let vRow = Tensor( - buffer: vOut.buffer, - offset: vOut.offset + r * kvDim * dtBytes, - shape: [kvDim], dtype: dt) - Ops.ropePartial( - qRow, position: startPosition + r, - headDim: headDim, rotaryDim: rotaryDim, - thetaBase: ropeTheta, on: cmd) - Ops.ropePartial( - kRow, position: startPosition + r, - headDim: headDim, rotaryDim: rotaryDim, - thetaBase: ropeTheta, on: cmd) - kRows.append(kRow.reshaped(to: [nKVHeads, headDim])) - vRows.append(vRow.reshaped(to: [nKVHeads, headDim])) - } - kv.appendRangeOnGPU(kRows: kRows, vRows: vRows, on: cmd) - - // ── SDPA — per-token loop over `sdpaDecode`. `Ops.sdpaMulti` - // would consolidate the T calls into one but it's head_dim-128 - // only today; Qwen3.6 attention is head_dim-256. The per-token - // sdpaDecode loop preserves correctness; the projection and - // norm wins above still amortise. A future head_dim-256 - // `sdpaMulti` variant (or transpose-to-prefill_mma) collapses - // these T launches into one. + // ── Q/K norm — one shared encoder over (T·nHeads) + (T·nKVHeads) + // rows. Collapses the two separate `rmsNormRows` dispatches into + // one `rmsNormRowsTwo` shared encoder, mirroring the single- + // token `forward` path. Saves 1 encoder begin/end per attn- + // layer prefill call. + let qNormed = Tensor.empty( + shape: queriesFlat.shape, dtype: queriesFlat.dtype, device: device) + let kNormed = Tensor.empty( + shape: kOut.shape, dtype: kOut.dtype, device: device) + Ops.rmsNormRowsTwo( + queriesFlat, weight: qNorm.weight, eps1: qNorm.eps, + nRows1: t * nHeads, rowSize1: headDim, into: qNormed, + kOut, weight: kNorm.weight, eps2: kNorm.eps, + nRows2: t * nKVHeads, rowSize2: headDim, into: kNormed, + on: cmd) + + // ── Batched RoPE over all T tokens on ONE shared encoder. + // `Ops.ropePartialManyTwo` pairs Q+K RoPE for each of the T + // rows in a single dispatch — saves T-1 encoder begin/end pairs + // per attn layer × 10 ≈ 5100 fewer dispatches at T=512. + let positions = device.makeBuffer(length: t * 4) + let positionsPtr = positions.contents().bindMemory( + to: UInt32.self, capacity: t) + for r in 0 ..< t { positionsPtr[r] = UInt32(startPosition + r) } + let positionsT = Tensor( + buffer: positions, offset: 0, shape: [t], dtype: .u32) + Ops.ropePartialManyTwo( + q: qNormed, qNHeads: nHeads, qRowStride: qDim, + k: kNormed, kNHeads: nKVHeads, kRowStride: kvDim, + positions: positionsT, t: t, + headDim: headDim, rotaryDim: rotaryDim, + thetaBase: ropeTheta, on: cmd) + + // ── KV append — batched K+V cache append. Reserves T + // sequential physical slots on host under lengthLock, then writes + // K/V for all T tokens via ONE shared encoder + 2 dispatches (vs + // the T-loop of `kvCacheUpdateKV` paired-encoder calls). + let appendPositions = Tensor.empty( + shape: [t], dtype: .u32, device: device) + kv.reserveSlotsManyOnHost(t: t, into: appendPositions) + kv.appendRangeOnGPUMany( + kFlat: kNormed, vFlat: vOut, t: t, + positions: appendPositions, on: cmd) + + // ── SDPA — fuse the T-token causal attention into ONE + // `Ops.sdpaMulti` dispatch. d=128 wraps `ffai_sdpa_multi`; + // d=256 (Qwen3.6-A3B full-attention layers) wraps + // `ffai_sdpa_multi_d256` — same semantics, head_dim-256 2-phase + // output reduction internally. T-loop `sdpaDecode` fallback only + // fires for other head_dims. let (cacheK, cacheV) = kv.prepareForAttention(on: cmd) - let attnAll = Tensor.empty(shape: [t * qDim], dtype: dt, device: device) - for r in 0 ..< t { - let qRow = Tensor( - buffer: qNormed.buffer, - offset: qNormed.offset + r * qDim * dtBytes, - shape: [nHeads, headDim], dtype: dt) - let outRow = Tensor( - buffer: attnAll.buffer, - offset: attnAll.offset + r * qDim * dtBytes, - shape: [nHeads, headDim], dtype: dt) - _ = Ops.sdpaDecode( - q: qRow, k: cacheK, v: cacheV, + let attnAll: Tensor + if headDim == 128 || headDim == 256 { + // qNormed is `[T, nHeads, headDim]` flat. Reshape into the + // shape sdpaMulti expects (same buffer, no copy). + let qBlock = qNormed.reshaped(to: [t, nHeads, headDim]) + attnAll = Ops.sdpaMulti( + q: qBlock, k: cacheK, v: cacheV, nQHeads: nHeads, nKVHeads: nKVHeads, headDim: headDim, - nKV: startPosition + r + 1, kvStride: kv.maxSeq, - scale: scale, on: cmd, into: outRow) + baseKV: startPosition, nQuery: t, kvStride: kv.maxSeq, + causal: true, scale: scale, on: cmd + ).reshaped(to: [t * qDim]) + } else { + let attnAllScratch = Tensor.empty( + shape: [t * qDim], dtype: dt, device: device) + for r in 0 ..< t { + let qRow = Tensor( + buffer: qNormed.buffer, + offset: qNormed.offset + r * qDim * dtBytes, + shape: [nHeads, headDim], dtype: dt) + let outRow = Tensor( + buffer: attnAllScratch.buffer, + offset: attnAllScratch.offset + r * qDim * dtBytes, + shape: [nHeads, headDim], dtype: dt) + _ = Ops.sdpaDecode( + q: qRow, k: cacheK, v: cacheV, + nQHeads: nHeads, nKVHeads: nKVHeads, headDim: headDim, + nKV: startPosition + r + 1, kvStride: kv.maxSeq, + scale: scale, on: cmd, into: outRow) + } + attnAll = attnAllScratch } - // ── Gated output * o_proj ──────────────────────────────────────── + // ── Gated output * o_proj — fused sigmoidMul on the prefill + // path (same kernel as the single-token forward path). var attnFlat = attnAll if let gateFlat { - attnFlat = Ops.mul(attnFlat, Ops.sigmoid(gateFlat, on: cmd), on: cmd) + attnFlat = Ops.sigmoidMul(attnFlat, gateFlat, on: cmd) } let attnRows = attnFlat.reshaped(to: [t, qDim]) return oProj.callMany(attnRows, t: t, on: cmd, device: device) @@ -1834,9 +2508,6 @@ public final class Qwen35GDNLayer: Module, DecoderLayer { "Qwen35GDNLayer.decodeMany: hFlat size \(hFlat.elementCount) ≠ T·hidden = \(t * hidden)" ) - let dt = hFlat.dtype - let dtBytes = dt.byteSize - // ── Pre-norm — one rmsNormRows over T rows ────────────────────── let xNormFlat = Ops.rmsNormRows( hFlat, weight: inputNorm.weight, eps: inputNorm.eps, @@ -1988,9 +2659,6 @@ public final class Qwen35AttentionLayer: Module, DecoderLayer { "Qwen35AttentionLayer.decodeMany: hFlat size \(hFlat.elementCount) " + "≠ T·hidden = \(t * hidden)") - let dt = hFlat.dtype - let dtBytes = dt.byteSize - // ── Pre-norm — one rmsNormRows over T rows ────────────────────── let xNormFlat = Ops.rmsNormRows( hFlat, weight: inputNorm.weight, eps: inputNorm.eps, @@ -2033,6 +2701,36 @@ public final class Qwen35AttentionLayer: Module, DecoderLayer { } } +// ─── finalNorm + lmHead fused dispatch helper ─────────────────────── +// +// Replaces the 2-dispatch chain `finalNorm(h)` + `lmHead(normed)` with +// a single `Ops.rmsNormQgemvInt4Fast` dispatch when the lmHead's +// underlying linear is a 4-bit `QuantizedLinear` that satisfies the +// fast kernel's constraints (in_dim multiple of 512, out_dim multiple +// of 8, group_size = 64). Falls back to the separate-chain form +// otherwise (dense lmHead, int8/int3 quant, smaller vocab, etc.). +private func qwen35FinalNormLmHead( + h: Tensor, finalNorm: RMSNorm, lmHead: AnyLinear, + on cmd: MTLCommandBuffer +) -> Tensor { + if let qLm = lmHead.inner as? QuantizedLinear, qLm.bits == 4, + qLm.groupSize == 64, + h.elementCount % 512 == 0, + qLm.weight.shape[0] % 8 == 0 + { + let outDim = qLm.weight.shape[0] + let logits = Tensor.empty(shape: [outDim], dtype: h.dtype) + Ops.rmsNormQgemvInt4Fast( + x: h, normWeight: finalNorm.weight, eps: finalNorm.eps, + qWeight: qLm.weight, qScales: qLm.scales, qBiases: qLm.biases, + on: cmd, into: logits) + return logits + } + // Fallback: separate dispatches (legacy two-step). + let normed = finalNorm(h, on: cmd) + return lmHead(normed, on: cmd) +} + // ─── Shared FFN helpers ────────────────────────────────────────────── /// Collect the `(name, tensor)` parameters of a layer's FFN half. @@ -2083,7 +2781,6 @@ private func qwen35ApplyFFNMany( let addCmd = device.makeCommandBuffer() let resultFlat = Ops.add(postMix, ffnOut, on: addCmd) addCmd.commit() - let _ = resultFlat return resultFlat } } @@ -2122,17 +2819,14 @@ private func qwen35ApplyFFN( } return result case .moe(let moe): - // Qwen35MoEFFN.forward commits `cmd`; run the residual add on a - // fresh buffer. We commit without waiting: the residual add - // tensor is the layer's output; the next layer queues onto a - // fresh workCmd and Metal hazard-tracks the read. - let ffnOut = moe.forward( + // Qwen35MoEFFN.forward commits `cmd`. Folds the post-FFN + // residual add into the final sigmoidScalarFMA via + // `sigmoidScalarFMAResidual` — when this is taken the returned + // tensor already includes `postMix + routed + sigmoid·shared`. + return moe.forward( ffnNorm, position: position, - cmd: cmd, device: device) - let addCmd = device.makeCommandBuffer() - let result = Ops.add(postMix, ffnOut, on: addCmd) - addCmd.commit() - return result + cmd: cmd, device: device, + residual: postMix) } } @@ -2305,8 +2999,10 @@ public final class Qwen35Model: LanguageModel { } // Final norm + lm_head queue onto the caller's pristine `cmd`. - let normed = finalNorm(h, on: cmd) - return lmHead(normed, on: cmd) + // Fused rmsNorm+qgemv when lmHead is 4-bit QuantizedLinear with + // matching constraints; falls back otherwise. + return qwen35FinalNormLmHead( + h: h, finalNorm: finalNorm, lmHead: lmHead, on: cmd) } /// LanguageModel-protocol entry point for chunked prefill. Delegates @@ -2392,6 +3088,31 @@ public final class Qwen35Model: LanguageModel { tokenIds: tokenIds, startPosition: startPosition, caches: caches, + returnAllLogits: false, + on: cmd, device: device) + } + + /// Multi-token forward returning logits at EVERY position — `[T, vocab]`. + /// Same KV/GDN-cache-writing semantics as `forwardMany`; differs only + /// in the output: instead of slicing the last row's hidden, applies + /// final RMSNorm + lm_head row-wise to all T positions. + /// + /// Spec-decode driver: a drafter proposes `γ` candidate tokens, the + /// target verifies via `forwardManyAllLogits(prefix + candidates)` to + /// get the per-position logits, then accepts up to the first reject. + public func forwardManyAllLogits( + tokenIds: [Int], startPosition: Int, + caches: [any LayerCacheProtocol], + on cmd: MTLCommandBuffer, device: Device + ) -> Tensor { + precondition( + !tokenIds.isEmpty, + "Qwen35Model.forwardManyAllLogits: tokenIds must not be empty") + return _forwardManyBatched( + tokenIds: tokenIds, + startPosition: startPosition, + caches: caches, + returnAllLogits: true, on: cmd, device: device) } @@ -2423,11 +3144,17 @@ public final class Qwen35Model: LanguageModel { /// Mixed-dispatch batched forwardMany. Attention layers fan to /// `decodeMany`; GDN (and any other) layer types stay per-token with /// a per-row `Ops.copy` blit-back into the running `[T, hidden]` - /// buffer. Embedding gathers all T tokens in one dispatch. Final - /// norm + lm_head run only on the LAST row, on the caller's `cmd`. + /// buffer. Embedding gathers all T tokens in one dispatch. + /// + /// Output shape depends on `returnAllLogits`: + /// * `false` (default, prefill driver): final norm + lm_head run + /// on the LAST row only → `[vocab]`. + /// * `true` (spec-decode verify): batched RMSNorm over T rows + + /// batched lm_head → `[T, vocab]`. private func _forwardManyBatched( tokenIds: [Int], startPosition: Int, caches: [any LayerCacheProtocol], + returnAllLogits: Bool, on cmd: MTLCommandBuffer, device: Device ) -> Tensor { let t = tokenIds.count @@ -2508,13 +3235,28 @@ public final class Qwen35Model: LanguageModel { workCmd.waitUntilCompleted() } - // ── Final norm + lm_head on the LAST row only ──────────────────── - let lastRow = Tensor( - buffer: h.buffer, - offset: h.offset + (t - 1) * hidden * dtBytes, - shape: [hidden], dtype: dt) - let normed = finalNorm(lastRow, on: cmd) - return lmHead(normed, on: cmd) + // ── Final norm + lm_head ───────────────────────────────────────── + // Two output shapes (see _forwardManyBatched docstring): + // * returnAllLogits == false: logits of the LAST row only, + // fused via qwen35FinalNormLmHead. + // * returnAllLogits == true: batched RMSNorm over T rows + + // batched lm_head — the spec-decode verify path needs + // per-position logits. + if !returnAllLogits { + let lastRow = Tensor( + buffer: h.buffer, + offset: h.offset + (t - 1) * hidden * dtBytes, + shape: [hidden], dtype: dt) + return qwen35FinalNormLmHead( + h: lastRow, finalNorm: finalNorm, lmHead: lmHead, on: cmd) + } + // All-T path: batched RMSNorm over T rows + batched lm_head. + let normedAll = Ops.rmsNormRows( + h, weight: finalNorm.weight, eps: finalNorm.eps, + nRows: t, rowSize: hidden, on: cmd) + return lmHead.callMany( + normedAll.reshaped(to: [t, hidden]), + t: t, on: cmd, device: device) } // ─── VLM embedding-input path ──────────────────────────────────── @@ -2564,8 +3306,10 @@ public final class Qwen35Model: LanguageModel { workCmd.waitUntilCompleted() } - let normed = finalNorm(h, on: cmd) - return lmHead(normed, on: cmd) + // Fused finalNorm+lmHead on the VLM embed-input forward path + // too. + return qwen35FinalNormLmHead( + h: h, finalNorm: finalNorm, lmHead: lmHead, on: cmd) } /// Raw embedding-table lookup for one text token. @@ -2713,6 +3457,55 @@ private func sliceHeadHalves35( return gathered.reshaped(to: [nHeads * headDim]) } +/// T-batched both-halves variant. Returns `(queriesFlat, gateFlat)` +/// — both `[T·nHeads·headDim]` flat — using a single shared-encoder +/// `Ops.gatherTwo` dispatch over the two interleaved row sets. +/// Replaces two back-to-back `sliceHeadHalvesMany35` calls in +/// `forwardMany` with one encoder begin/end pair. +private func sliceHeadHalvesManyBoth35( + _ q2T: Tensor, t: Int, + nHeads: Int, headDim: Int, + on cmd: MTLCommandBuffer, + device: Device +) -> (Tensor, Tensor) { + precondition( + q2T.elementCount == t * nHeads * 2 * headDim, + "sliceHeadHalvesManyBoth35: q2T must be [T, nHeads, 2·headDim]") + let table = q2T.reshaped(to: [t * nHeads * 2, headDim]) + let nIdx = t * nHeads + var rowsFirst = [UInt32](repeating: 0, count: nIdx) + var rowsSecond = [UInt32](repeating: 0, count: nIdx) + for r in 0 ..< t { + let base = UInt32(r * 2 * nHeads) + let rowBase = r * nHeads + for h in 0 ..< nHeads { + rowsFirst[rowBase + h] = base + UInt32(2 * h) + rowsSecond[rowBase + h] = base + UInt32(2 * h) + 1 + } + } + let firstBuf = device.makeBuffer(length: nIdx * 4) + rowsFirst.withUnsafeBytes { + _ = memcpy(firstBuf.contents(), $0.baseAddress!, nIdx * 4) + } + let secondBuf = device.makeBuffer(length: nIdx * 4) + rowsSecond.withUnsafeBytes { + _ = memcpy(secondBuf.contents(), $0.baseAddress!, nIdx * 4) + } + let idxFirst = Tensor(buffer: firstBuf, offset: 0, shape: [nIdx], dtype: .u32) + let idxSecond = Tensor(buffer: secondBuf, offset: 0, shape: [nIdx], dtype: .u32) + let out1 = Tensor.empty(shape: [nIdx, headDim], dtype: q2T.dtype, device: device) + let out2 = Tensor.empty(shape: [nIdx, headDim], dtype: q2T.dtype, device: device) + Ops.gatherTwo( + table: table, + ids1: idxFirst, into: out1, + ids2: idxSecond, into: out2, + on: cmd) + return ( + out1.reshaped(to: [nIdx * headDim]), + out2.reshaped(to: [nIdx * headDim]) + ) +} + /// T-batched variant of `sliceHeadHalves35`. Input `q2T` is /// `[T, nHeads, 2·headDim]` — each of T rows has its own /// `[nHeads, 2·headDim]` block. Returns `[T · nHeads · headDim]` flat diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift index ec4695be..cd298ca0 100644 --- a/Sources/FFAI/Ops/Ops.swift +++ b/Sources/FFAI/Ops/Ops.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -345,6 +345,54 @@ public enum Ops { return result } + /// Two embedding lookups against the SAME `table` on one compute + /// encoder. Used when two parallel id streams need to read the + /// same embedding (or, in Qwen3 attention, when a single linear + /// projection result is split into two head-half slices via + /// `ids = [0..nHeads]` + `ids = [nHeads..2·nHeads]`). The table + /// binding + `dim` constant are set once and the per-call + /// `(ids, out)` pair rotates per dispatch. + public static func gatherTwo( + table: Tensor, + ids1: Tensor, into out1: Tensor, + ids2: Tensor, into out2: Tensor, + on cmd: MTLCommandBuffer + ) { + precondition(table.shape.count == 2, "Ops.gatherTwo: table must be 2D") + precondition( + ids1.dtype == .u32 && ids2.dtype == .u32, + "Ops.gatherTwo: ids must be u32") + precondition( + out1.dtype == table.dtype && out2.dtype == table.dtype, + "Ops.gatherTwo: out dtype must match table") + let dim = table.shape[1] + let psoName: String + switch table.dtype { + case .f32: psoName = "ffai_gather_f32" + case .f16: psoName = "ffai_gather_f16" + case .bf16: psoName = "ffai_gather_bf16" + default: fatalError("Ops.gatherTwo: unsupported dtype \(table.dtype)") + } + let pso = PSOCache.shared.pipelineState(for: psoName) + guard let enc = cmd.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(pso) + var dimV = UInt32(dim) + enc.setBytes(&dimV, length: 4, index: 3) + enc.setBuffer(table.buffer, offset: table.offset, index: 0) + @inline(__always) + func dispatch(_ ids: Tensor, _ out: Tensor) { + let n = ids.elementCount + let totalThreads = n * dim + let (grid, tg) = elementwiseGrid(totalThreads) + enc.setBuffer(ids.buffer, offset: ids.offset, index: 1) + enc.setBuffer(out.buffer, offset: out.offset, index: 2) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + } + dispatch(ids1, out1) + dispatch(ids2, out2) + enc.endEncoding() + } + /// Cooperative-thread matrix-vector multiply. weight: [out_dim, in_dim], /// input: [in_dim], output: [out_dim]. One threadgroup per output row; /// threads cooperate on the dot-product reduction. @@ -514,6 +562,99 @@ public enum Ops { return result } + /// Two `rmsNormRows` dispatches in ONE compute encoder. Used by + /// the Qwen3 attention mixer's pre-RoPE Q-norm + K-norm pair: both + /// run the fast `mt_rms_norm_*` kernel (TPG = rowSize / 4), share + /// the same eps when set from the same model config, and there's + /// no data dependency between them. + /// + /// `rowSize1` and `rowSize2` must both clear the fast-path bounds + /// (`≤ 4096` and a multiple of 128 — checked via + /// `OpsValidation.validateRmsNorm`). For wider rows the caller + /// must fall back to two `rmsNormRows` calls so the wide kernel is + /// picked individually. + /// + /// `eps` is forwarded via 1-element f32 buffers. If both eps + /// values are equal we share a single alloc, otherwise two are + /// allocated. + public static func rmsNormRowsTwo( + _ x1: Tensor, weight w1: Tensor, eps1: Float, + nRows1: Int, rowSize1: Int, into out1: Tensor, + _ x2: Tensor, weight w2: Tensor, eps2: Float, + nRows2: Int, rowSize2: Int, into out2: Tensor, + on cmd: MTLCommandBuffer + ) { + precondition( + x1.dtype == w1.dtype && x2.dtype == w2.dtype && x1.dtype == x2.dtype, + "Ops.rmsNormRowsTwo: dtype mismatch") + precondition( + x1.elementCount == nRows1 * rowSize1, + "Ops.rmsNormRowsTwo: x1 size \(x1.elementCount) ≠ nRows1·rowSize1") + precondition( + x2.elementCount == nRows2 * rowSize2, + "Ops.rmsNormRowsTwo: x2 size \(x2.elementCount) ≠ nRows2·rowSize2") + precondition( + w1.elementCount == rowSize1 && w2.elementCount == rowSize2, + "Ops.rmsNormRowsTwo: weight size must equal rowSize") + precondition( + out1.elementCount == x1.elementCount && out2.elementCount == x2.elementCount, + "Ops.rmsNormRowsTwo: output element-count must match input") + precondition( + out1.dtype == x1.dtype && out2.dtype == x2.dtype, + "Ops.rmsNormRowsTwo: output dtype must match input") + if let reason = OpsValidation.validateRmsNorm(n: rowSize1) { + preconditionFailure("Ops.rmsNormRowsTwo (#1): \(reason)") + } + if let reason = OpsValidation.validateRmsNorm(n: rowSize2) { + preconditionFailure("Ops.rmsNormRowsTwo (#2): \(reason)") + } + let psoName: String + switch x1.dtype { + case .f32: psoName = "mt_rms_norm_f32" + case .f16: psoName = "mt_rms_norm_f16" + case .bf16: psoName = "mt_rms_norm_bf16" + default: fatalError("Ops.rmsNormRowsTwo: unsupported dtype \(x1.dtype)") + } + let pso = PSOCache.shared.pipelineState(for: psoName) + guard let enc = cmd.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(pso) + // Allocate one or two eps buffers depending on whether the two + // eps values match — RMSNorm pairs typically come from the + // same checkpoint config and so share a single eps in practice. + let epsBuf1: MTLBuffer = { + let b = device.makeBuffer(length: 4) + var v = eps1 + memcpy(b.contents(), &v, 4) + return b + }() + let epsBuf2: MTLBuffer = { + if eps1 == eps2 { return epsBuf1 } + let b = device.makeBuffer(length: 4) + var v = eps2 + memcpy(b.contents(), &v, 4) + return b + }() + @inline(__always) + func dispatch( + _ x: Tensor, _ w: Tensor, _ out: Tensor, + _ epsBuf: MTLBuffer, _ n: Int, _ nRows: Int + ) { + let tgWidth = n / 4 + let grid = MTLSize(width: nRows * tgWidth, height: 1, depth: 1) + let tg = MTLSize(width: tgWidth, height: 1, depth: 1) + enc.setBuffer(x.buffer, offset: x.offset, index: 0) + enc.setBuffer(w.buffer, offset: w.offset, index: 1) + enc.setBuffer(out.buffer, offset: out.offset, index: 2) + enc.setBuffer(epsBuf, offset: 0, index: 3) + var nU = UInt32(n) + enc.setBytes(&nU, length: 4, index: 4) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + } + dispatch(x1, w1, out1, epsBuf1, rowSize1, nRows1) + dispatch(x2, w2, out2, epsBuf2, rowSize2, nRows2) + enc.endEncoding() + } + /// Fused (residual `a + b`) + RMSNorm. Returns **both** outputs: /// `residual` = `a + b` (the next layer's residual stream) and /// `normed` = `RMSNorm(a + b, weight)` (the input to the next @@ -840,6 +981,197 @@ public enum Ops { } } + /// Partial RoPE rotation on TWO tensors (typically Q + K) in ONE + /// compute encoder. Both tensors share the same `(headDim, + /// rotaryDim, position, thetaBase, scaling)` and dtype. The + /// kernel writes in-place exactly like `ropePartial`, so passing + /// `out == qk` is required. + public static func ropePartialTwo( + _ q: Tensor, _ k: Tensor, position: Int, + headDim: Int, rotaryDim: Int, + thetaBase: Float, + scaling: RoPEScaling = .none, + on cmd: MTLCommandBuffer + ) { + precondition( + q.dtype == k.dtype, + "Ops.ropePartialTwo: dtype mismatch") + precondition( + q.elementCount % headDim == 0 && k.elementCount % headDim == 0, + "Ops.ropePartialTwo: sizes must be multiples of headDim") + precondition( + rotaryDim > 0 && rotaryDim <= headDim && rotaryDim % 2 == 0, + "Ops.ropePartialTwo: rotaryDim (\(rotaryDim)) must be even and in 1...headDim") + let psoName: String + switch q.dtype { + case .f32: psoName = "ffai_rope_llama_f32" + case .f16: psoName = "ffai_rope_llama_f16" + case .bf16: psoName = "ffai_rope_llama_bf16" + default: fatalError("Ops.ropePartialTwo: unsupported dtype \(q.dtype)") + } + let halfRotary = rotaryDim / 2 + let pso = PSOCache.shared.pipelineState(for: psoName) + guard let enc = cmd.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(pso) + // RoPE constants are shared across q and k — set ONCE. + var hd = UInt32(headDim) + var half = UInt32(halfRotary) + var pos = UInt32(position) + var theta = thetaBase + var scaleFactor = scaling.scaleFactor + var lowFreq = scaling.lowFreqFactor + var highFreq = scaling.highFreqFactor + var origMax = scaling.originalMaxPosition + enc.setBytes(&hd, length: 4, index: 2) + enc.setBytes(&half, length: 4, index: 3) + enc.setBytes(&pos, length: 4, index: 4) + enc.setBytes(&theta, length: 4, index: 5) + enc.setBytes(&scaleFactor, length: 4, index: 6) + enc.setBytes(&lowFreq, length: 4, index: 7) + enc.setBytes(&highFreq, length: 4, index: 8) + enc.setBytes(&origMax, length: 4, index: 9) + let tg = MTLSize(width: 1, height: 1, depth: 1) + @inline(__always) + func dispatch(_ t: Tensor) { + let nHeads = t.elementCount / headDim + let grid = MTLSize(width: nHeads, height: halfRotary, depth: 1) + enc.setBuffer(t.buffer, offset: t.offset, index: 0) + enc.setBuffer(t.buffer, offset: t.offset, index: 1) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + } + dispatch(q) + dispatch(k) + enc.endEncoding() + } + + /// Batched per-row RoPE: applies position-dependent rotation to T + /// rows of `qk` in ONE dispatch. `positions` is a `[T]` u32 buffer + /// (per-row physical slot index). In-place rotation. Saves T-1 + /// encoder begin/end pairs versus a per-row `ropePartial` T-loop; + /// at Qwen3.6-A3B prefill T=512 × 10 attn layers that's ≈ 5100 + /// fewer dispatches per prefill call. + public static func ropePartialMany( + _ qk: Tensor, positions: Tensor, + t: Int, nHeads: Int, + headDim: Int, rotaryDim: Int, + rowStride: Int, + thetaBase: Float, + scaling: RoPEScaling = .none, + on cmd: MTLCommandBuffer + ) { + precondition( + positions.dtype == .u32, + "Ops.ropePartialMany: positions must be .u32") + precondition( + positions.elementCount == t, + "Ops.ropePartialMany: positions count must equal T") + precondition( + qk.elementCount == t * rowStride, + "Ops.ropePartialMany: qk size \(qk.elementCount) ≠ T·rowStride = \(t * rowStride)") + precondition( + rotaryDim > 0 && rotaryDim <= headDim && rotaryDim % 2 == 0, + "Ops.ropePartialMany: rotaryDim must be even and in 1...headDim") + let halfRotary = rotaryDim / 2 + let psoName: String + switch qk.dtype { + case .f32: psoName = "ffai_rope_llama_many_f32" + case .f16: psoName = "ffai_rope_llama_many_f16" + case .bf16: psoName = "ffai_rope_llama_many_bf16" + default: fatalError("Ops.ropePartialMany: unsupported dtype \(qk.dtype)") + } + let pso = PSOCache.shared.pipelineState(for: psoName) + guard let enc = cmd.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(pso) + enc.setBuffer(qk.buffer, offset: qk.offset, index: 0) + enc.setBuffer(positions.buffer, offset: positions.offset, index: 1) + enc.setBuffer(qk.buffer, offset: qk.offset, index: 2) + var hd = UInt32(headDim) + var half = UInt32(halfRotary) + var stride = UInt32(rowStride) + var theta = thetaBase + var sf = scaling.scaleFactor + var lf = scaling.lowFreqFactor + var hf = scaling.highFreqFactor + var om = scaling.originalMaxPosition + enc.setBytes(&hd, length: 4, index: 3) + enc.setBytes(&half, length: 4, index: 4) + enc.setBytes(&stride, length: 4, index: 5) + enc.setBytes(&theta, length: 4, index: 6) + enc.setBytes(&sf, length: 4, index: 7) + enc.setBytes(&lf, length: 4, index: 8) + enc.setBytes(&hf, length: 4, index: 9) + enc.setBytes(&om, length: 4, index: 10) + let grid = MTLSize(width: t, height: nHeads, depth: halfRotary) + let tg = MTLSize(width: 1, height: 1, depth: 1) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + } + + /// Pair of `ropePartialMany` calls (Q + K) sharing ONE compute + /// encoder. Both buffers consume the same `positions`, `headDim`, + /// `rotaryDim`, `thetaBase`, and `scaling`; only the per-buffer + /// `nHeads` × `rowStride` differs. + public static func ropePartialManyTwo( + q: Tensor, qNHeads: Int, qRowStride: Int, + k: Tensor, kNHeads: Int, kRowStride: Int, + positions: Tensor, t: Int, + headDim: Int, rotaryDim: Int, + thetaBase: Float, + scaling: RoPEScaling = .none, + on cmd: MTLCommandBuffer + ) { + precondition(q.dtype == k.dtype, "Ops.ropePartialManyTwo: dtype mismatch") + precondition( + positions.dtype == .u32, + "Ops.ropePartialManyTwo: positions must be .u32") + precondition( + positions.elementCount == t, + "Ops.ropePartialManyTwo: positions count must equal T") + precondition( + rotaryDim > 0 && rotaryDim <= headDim && rotaryDim % 2 == 0, + "Ops.ropePartialManyTwo: rotaryDim must be even and in 1...headDim") + let halfRotary = rotaryDim / 2 + let psoName: String + switch q.dtype { + case .f32: psoName = "ffai_rope_llama_many_f32" + case .f16: psoName = "ffai_rope_llama_many_f16" + case .bf16: psoName = "ffai_rope_llama_many_bf16" + default: fatalError("Ops.ropePartialManyTwo: unsupported dtype \(q.dtype)") + } + let pso = PSOCache.shared.pipelineState(for: psoName) + guard let enc = cmd.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(pso) + // Shared bindings: positions + constants set ONCE. + enc.setBuffer(positions.buffer, offset: positions.offset, index: 1) + var hd = UInt32(headDim) + var half = UInt32(halfRotary) + var theta = thetaBase + var sf = scaling.scaleFactor + var lf = scaling.lowFreqFactor + var hf = scaling.highFreqFactor + var om = scaling.originalMaxPosition + enc.setBytes(&hd, length: 4, index: 3) + enc.setBytes(&half, length: 4, index: 4) + enc.setBytes(&theta, length: 4, index: 6) + enc.setBytes(&sf, length: 4, index: 7) + enc.setBytes(&lf, length: 4, index: 8) + enc.setBytes(&hf, length: 4, index: 9) + enc.setBytes(&om, length: 4, index: 10) + let tg = MTLSize(width: 1, height: 1, depth: 1) + @inline(__always) + func dispatch(_ buf: Tensor, _ nHeads: Int, _ rowStride: Int) { + var stride = UInt32(rowStride) + enc.setBuffer(buf.buffer, offset: buf.offset, index: 0) + enc.setBuffer(buf.buffer, offset: buf.offset, index: 2) + enc.setBytes(&stride, length: 4, index: 5) + let grid = MTLSize(width: t, height: nHeads, depth: halfRotary) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + } + dispatch(q, qNHeads, qRowStride) + dispatch(k, kNHeads, kRowStride) + enc.endEncoding() + } + /// YaRN RoPE parameters. `low` / `high` are the correction-range /// bounds — precomputed via `RoPEYaRN.from(...)` since they need a /// `floor`/`ceil`/`ln` computation that is constant across the @@ -1319,156 +1651,1184 @@ public enum Ops { on: cmd, into: out) } - /// GPU argmax over a 1D logits tensor. Caller supplies a 1-element - /// u32 output buffer. Uses the cooperative 256-thread Reduction - /// kernel — one threadgroup, ~80-300 KB / vocab logits in registers. - /// Mamba 2 / Mamba 1D depthwise causal-conv step — streaming - /// decode form. One thread per channel. `state` is the rolling - /// window of the last `kernelSize - 1` inputs (shape - /// `[kernelSize - 1, nChannels]`); shifted in-place after compute. - /// Activation (Mamba 2 follows the conv with SiLU) is the caller's - /// concern — kept separate for composability. - public static func conv1dCausalStep( - x: Tensor, w: Tensor, b: Tensor, - state: Tensor, into y: Tensor, - nChannels: Int, kernelSize: Int, + /// Batched int4 dequantGemv on TWO projections sharing one input + /// in ONE compute encoder. Used by `Qwen35MoEFFN.forward` for the + /// per-expert gate + up pair: both projections read the same + /// `[hidden]` post-norm activation, so we set the `input` binding + /// once and rotate `(weight, scales, biases, output)` per dispatch. + /// Saves one encoder begin/end pair per call. + /// + /// All outputs must share dtype with `input`; `groupSize` and the + /// 4-bit packing are identical across the two projections. + public static func dequantGemvInt4Two( + input: Tensor, + w0: Tensor, s0: Tensor, b0: Tensor, out0: Tensor, + w1: Tensor, s1: Tensor, b1: Tensor, out1: Tensor, + groupSize: Int = 64, on cmd: MTLCommandBuffer ) { precondition( - w.dtype == x.dtype && b.dtype == x.dtype - && state.dtype == x.dtype && y.dtype == x.dtype, - "Ops.conv1dCausalStep: every tensor must share dtype") - let grid = MTLSize(width: nChannels, height: 1, depth: 1) - let tg = MTLSize(width: 1, height: 1, depth: 1) + input.dtype == out0.dtype && input.dtype == out1.dtype, + "Ops.dequantGemvInt4Two: dtype mismatch") + let psoName: String + switch input.dtype { + case .f32: psoName = "dequant_gemv_int4_f32" + case .f16: psoName = "dequant_gemv_int4_f16" + case .bf16: psoName = "dequant_gemv_int4_bf16" + default: fatalError("Ops.dequantGemvInt4Two: unsupported dtype \(input.dtype)") + } + let pso = PSOCache.shared.pipelineState(for: psoName) + guard let enc = cmd.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(pso) + enc.setBuffer(input.buffer, offset: input.offset, index: 3) + // `inDim` derives from the packed-row width: 4-bit packs 8 + // weights per u32 word and `weight.shape[1]` counts words. + let packedPerRow = w0.shape[1] + let inDim = packedPerRow * 32 / 4 + var inDimV = UInt32(inDim) + var groupSizeV = UInt32(groupSize) + enc.setBytes(&inDimV, length: 4, index: 5) + enc.setBytes(&groupSizeV, length: 4, index: 6) + let tgWidth = 256 + let tg = MTLSize(width: tgWidth, height: 1, depth: 1) + @inline(__always) + func dispatch(_ w: Tensor, _ s: Tensor, _ b: Tensor, _ out: Tensor) { + enc.setBuffer(w.buffer, offset: w.offset, index: 0) + enc.setBuffer(s.buffer, offset: s.offset, index: 1) + enc.setBuffer(b.buffer, offset: b.offset, index: 2) + enc.setBuffer(out.buffer, offset: out.offset, index: 4) + let outDim = w.shape[0] + let grid = MTLSize(width: outDim * tgWidth, height: 1, depth: 1) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + } + dispatch(w0, s0, b0, out0) + dispatch(w1, s1, b1, out1) + enc.endEncoding() + } + + /// Batched int4 dequantGemv on THREE projections sharing one input + /// in ONE compute encoder. Used by the Qwen3.5/3.6 attention mixer + /// for the q/k/v projection triplet. + public static func dequantGemvInt4Three( + input: Tensor, + w0: Tensor, s0: Tensor, b0: Tensor, out0: Tensor, + w1: Tensor, s1: Tensor, b1: Tensor, out1: Tensor, + w2: Tensor, s2: Tensor, b2: Tensor, out2: Tensor, + groupSize: Int = 64, + on cmd: MTLCommandBuffer + ) { + precondition( + input.dtype == out0.dtype, + "Ops.dequantGemvInt4Three: dtype mismatch") + let psoName: String + switch input.dtype { + case .f32: psoName = "dequant_gemv_int4_f32" + case .f16: psoName = "dequant_gemv_int4_f16" + case .bf16: psoName = "dequant_gemv_int4_bf16" + default: fatalError("Ops.dequantGemvInt4Three: unsupported dtype \(input.dtype)") + } + let pso = PSOCache.shared.pipelineState(for: psoName) + guard let enc = cmd.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(pso) + enc.setBuffer(input.buffer, offset: input.offset, index: 3) + let packedPerRow = w0.shape[1] + let inDim = packedPerRow * 32 / 4 + var inDimV = UInt32(inDim) + var groupSizeV = UInt32(groupSize) + enc.setBytes(&inDimV, length: 4, index: 5) + enc.setBytes(&groupSizeV, length: 4, index: 6) + let tgWidth = 256 + let tg = MTLSize(width: tgWidth, height: 1, depth: 1) + @inline(__always) + func dispatch(_ w: Tensor, _ s: Tensor, _ b: Tensor, _ out: Tensor) { + enc.setBuffer(w.buffer, offset: w.offset, index: 0) + enc.setBuffer(s.buffer, offset: s.offset, index: 1) + enc.setBuffer(b.buffer, offset: b.offset, index: 2) + enc.setBuffer(out.buffer, offset: out.offset, index: 4) + let outDim = w.shape[0] + let grid = MTLSize(width: outDim * tgWidth, height: 1, depth: 1) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + } + dispatch(w0, s0, b0, out0) + dispatch(w1, s1, b1, out1) + dispatch(w2, s2, b2, out2) + enc.endEncoding() + } + + /// Batched int4 dequantGemv on FOUR projections sharing one input + /// in ONE compute encoder. Used by the Qwen3.5/3.6 GDN mixer where + /// the four input projections (qkv, z, b, a) all read the same + /// xNorm output. + public static func dequantGemvInt4Four( + input: Tensor, + w0: Tensor, s0: Tensor, b0: Tensor, out0: Tensor, + w1: Tensor, s1: Tensor, b1: Tensor, out1: Tensor, + w2: Tensor, s2: Tensor, b2: Tensor, out2: Tensor, + w3: Tensor, s3: Tensor, b3: Tensor, out3: Tensor, + groupSize: Int = 64, + on cmd: MTLCommandBuffer + ) { + precondition( + input.dtype == out0.dtype, + "Ops.dequantGemvInt4Four: dtype mismatch") + let psoName: String + switch input.dtype { + case .f32: psoName = "dequant_gemv_int4_f32" + case .f16: psoName = "dequant_gemv_int4_f16" + case .bf16: psoName = "dequant_gemv_int4_bf16" + default: fatalError("Ops.dequantGemvInt4Four: unsupported dtype \(input.dtype)") + } + let pso = PSOCache.shared.pipelineState(for: psoName) + guard let enc = cmd.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(pso) + enc.setBuffer(input.buffer, offset: input.offset, index: 3) + let packedPerRow = w0.shape[1] + let inDim = packedPerRow * 32 / 4 + var inDimV = UInt32(inDim) + var groupSizeV = UInt32(groupSize) + enc.setBytes(&inDimV, length: 4, index: 5) + enc.setBytes(&groupSizeV, length: 4, index: 6) + let tgWidth = 256 + let tg = MTLSize(width: tgWidth, height: 1, depth: 1) + @inline(__always) + func dispatch(_ w: Tensor, _ s: Tensor, _ b: Tensor, _ out: Tensor) { + enc.setBuffer(w.buffer, offset: w.offset, index: 0) + enc.setBuffer(s.buffer, offset: s.offset, index: 1) + enc.setBuffer(b.buffer, offset: b.offset, index: 2) + enc.setBuffer(out.buffer, offset: out.offset, index: 4) + let outDim = w.shape[0] + let grid = MTLSize(width: outDim * tgWidth, height: 1, depth: 1) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + } + dispatch(w0, s0, b0, out0) + dispatch(w1, s1, b1, out1) + dispatch(w2, s2, b2, out2) + dispatch(w3, s3, b3, out3) + enc.endEncoding() + } + + /// Batched int4 dequant-GEMV on N projections with DIFFERENT + /// inputs sharing ONE compute encoder. Unlike `dequantGemvInt4Two` + /// / `Three` / `Four` (which all share a single input), the per- + /// expert MoE down phase has N independent activations — one per + /// chosen expert — but all projections share PSO, `inDim`, and + /// `groupSize`. The wrapper sets the constexprs once and rotates + /// `(weight, scales, biases, input, output)` per dispatch. Saves + /// N-1 encoder begin/end pairs versus N independent + /// `dequantGemvInt4` calls. + public static func dequantGemvInt4Many( + weights: [Tensor], scales: [Tensor], biases: [Tensor], + inputs: [Tensor], outputs: [Tensor], + groupSize: Int = 64, + on cmd: MTLCommandBuffer + ) { + let n = weights.count + precondition( + scales.count == n && biases.count == n && inputs.count == n && outputs.count == n, + "Ops.dequantGemvInt4Many: count mismatch") + guard n > 0 else { return } + let dtype = inputs[0].dtype + let psoName: String + switch dtype { + case .f32: psoName = "dequant_gemv_int4_f32" + case .f16: psoName = "dequant_gemv_int4_f16" + case .bf16: psoName = "dequant_gemv_int4_bf16" + default: fatalError("Ops.dequantGemvInt4Many: unsupported dtype \(dtype)") + } + let pso = PSOCache.shared.pipelineState(for: psoName) + guard let enc = cmd.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(pso) + let packedPerRow = weights[0].shape[1] + let inDim = packedPerRow * 32 / 4 // bits = 4 → 8 weights / u32 + var inDimV = UInt32(inDim) + var groupSizeV = UInt32(groupSize) + enc.setBytes(&inDimV, length: 4, index: 5) + enc.setBytes(&groupSizeV, length: 4, index: 6) + let tgWidth = 256 + let tg = MTLSize(width: tgWidth, height: 1, depth: 1) + for i in 0 ..< n { + precondition( + weights[i].shape[1] == packedPerRow, + "Ops.dequantGemvInt4Many: inDim varies at index \(i)") + precondition( + inputs[i].dtype == dtype && outputs[i].dtype == dtype, + "Ops.dequantGemvInt4Many: dtype varies at index \(i)") + enc.setBuffer(weights[i].buffer, offset: weights[i].offset, index: 0) + enc.setBuffer(scales[i].buffer, offset: scales[i].offset, index: 1) + enc.setBuffer(biases[i].buffer, offset: biases[i].offset, index: 2) + enc.setBuffer(inputs[i].buffer, offset: inputs[i].offset, index: 3) + enc.setBuffer(outputs[i].buffer, offset: outputs[i].offset, index: 4) + let outDim = weights[i].shape[0] + let grid = MTLSize(width: outDim * tgWidth, height: 1, depth: 1) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + } + enc.endEncoding() + } + + /// Fused Q/K/V int4 dequant-GEMV in ONE dispatch via + /// `ffai_batched_qkv_qgemv_fast`. Replaces the 3-dispatch + /// `dequantGemvInt4Three` shared-encoder form: the matrix is + /// selected by `program_id<2>()` across a `[ceil(maxOut/8), 1, 3]` + /// threadgroup grid, so all 3 projections share a single kernel + /// launch and read `x` once into TG memory instead of three times + /// from DRAM. Saves 2 encoder begin/end pairs per attention layer + /// at decode T=1. + /// + /// Output layout: q, k, v concatenated in that order in `out`; + /// caller slices into the three regions. + /// + /// Kernel constraints (mirror `ffai_rms_norm_qgemv_fast`): + /// - `in_dim` MUST be a multiple of 512. + /// - Each of `out_q`, `out_k`, `out_v` MUST be a multiple of 8. + /// - `group_size` MUST be 64. + /// - TPG = 64. + public static func batchedQkvQgemvInt4Fast( + x: Tensor, + wQ: Tensor, scalesQ: Tensor, biasesQ: Tensor, + wK: Tensor, scalesK: Tensor, biasesK: Tensor, + wV: Tensor, scalesV: Tensor, biasesV: Tensor, + outQ: Int, outK: Int, outV: Int, + on cmd: MTLCommandBuffer, + into out: Tensor + ) { + precondition( + wQ.dtype == .u32 && wK.dtype == .u32 && wV.dtype == .u32, + "Ops.batchedQkvQgemvInt4Fast: w_* must be u32-packed") + let packedPerRow = wQ.shape[1] + let inDim = packedPerRow * 8 + // K and V share the same in_dim as Q — the kernel reads ONE x row + // and walks all three weight tensors at that packed width. A + // mis-packed wK/wV would silently miscompute the K/V projections. + precondition( + wK.shape[1] == packedPerRow && wV.shape[1] == packedPerRow, + "Ops.batchedQkvQgemvInt4Fast: wK/wV packed width (\(wK.shape[1])/\(wV.shape[1])) must equal wQ (\(packedPerRow))" + ) + precondition( + x.elementCount == inDim, + "Ops.batchedQkvQgemvInt4Fast: x.elementCount \(x.elementCount) ≠ inDim \(inDim)") + precondition( + out.elementCount == outQ + outK + outV, + "Ops.batchedQkvQgemvInt4Fast: out.elementCount must be q+k+v") + precondition( + inDim % 512 == 0, + "Ops.batchedQkvQgemvInt4Fast: in_dim must be a multiple of 512") + precondition( + outQ % 8 == 0 && outK % 8 == 0 && outV % 8 == 0, + "Ops.batchedQkvQgemvInt4Fast: out_q/k/v must each be a multiple of 8") + let groupSize = 64 + let maxOut = max(outQ, max(outK, outV)) + // Grid: [ceil(maxOut/8) * TPG, 1, 3] — TG count per matrix axis + // times TPG=64 lanes; 3 matrices. + let tpg = 64 + let nTiles = maxOut / 8 + let grid = MTLSize(width: nTiles * tpg, height: 1, depth: 3) + let tg = MTLSize(width: tpg, height: 1, depth: 1) switch x.dtype { case .f32: - MetalTileKernels.conv1d_causal_step_f32( + MetalTileKernels.ffai_batched_qkv_qgemv_fast_f32( x: x.buffer, xOffset: x.offset, - w: w.buffer, wOffset: w.offset, - b: b.buffer, bOffset: b.offset, - state: state.buffer, stateOffset: state.offset, - y: y.buffer, yOffset: y.offset, - n_channels: UInt32(nChannels), kernel_size: UInt32(kernelSize), + w_q: wQ.buffer, w_qOffset: wQ.offset, + scales_q: scalesQ.buffer, scales_qOffset: scalesQ.offset, + biases_q: biasesQ.buffer, biases_qOffset: biasesQ.offset, + w_k: wK.buffer, w_kOffset: wK.offset, + scales_k: scalesK.buffer, scales_kOffset: scalesK.offset, + biases_k: biasesK.buffer, biases_kOffset: biasesK.offset, + w_v: wV.buffer, w_vOffset: wV.offset, + scales_v: scalesV.buffer, scales_vOffset: scalesV.offset, + biases_v: biasesV.buffer, biases_vOffset: biasesV.offset, + out: out.buffer, outOffset: out.offset, + out_q: UInt32(outQ), out_k: UInt32(outK), out_v: UInt32(outV), + in_dim: UInt32(inDim), group_size: UInt32(groupSize), gridSize: grid, threadgroupSize: tg, on: cmd) case .f16: - MetalTileKernels.conv1d_causal_step_f16( + MetalTileKernels.ffai_batched_qkv_qgemv_fast_f16( x: x.buffer, xOffset: x.offset, - w: w.buffer, wOffset: w.offset, - b: b.buffer, bOffset: b.offset, - state: state.buffer, stateOffset: state.offset, - y: y.buffer, yOffset: y.offset, - n_channels: UInt32(nChannels), kernel_size: UInt32(kernelSize), + w_q: wQ.buffer, w_qOffset: wQ.offset, + scales_q: scalesQ.buffer, scales_qOffset: scalesQ.offset, + biases_q: biasesQ.buffer, biases_qOffset: biasesQ.offset, + w_k: wK.buffer, w_kOffset: wK.offset, + scales_k: scalesK.buffer, scales_kOffset: scalesK.offset, + biases_k: biasesK.buffer, biases_kOffset: biasesK.offset, + w_v: wV.buffer, w_vOffset: wV.offset, + scales_v: scalesV.buffer, scales_vOffset: scalesV.offset, + biases_v: biasesV.buffer, biases_vOffset: biasesV.offset, + out: out.buffer, outOffset: out.offset, + out_q: UInt32(outQ), out_k: UInt32(outK), out_v: UInt32(outV), + in_dim: UInt32(inDim), group_size: UInt32(groupSize), gridSize: grid, threadgroupSize: tg, on: cmd) case .bf16: - MetalTileKernels.conv1d_causal_step_bf16( + MetalTileKernels.ffai_batched_qkv_qgemv_fast_bf16( x: x.buffer, xOffset: x.offset, - w: w.buffer, wOffset: w.offset, - b: b.buffer, bOffset: b.offset, - state: state.buffer, stateOffset: state.offset, - y: y.buffer, yOffset: y.offset, - n_channels: UInt32(nChannels), kernel_size: UInt32(kernelSize), + w_q: wQ.buffer, w_qOffset: wQ.offset, + scales_q: scalesQ.buffer, scales_qOffset: scalesQ.offset, + biases_q: biasesQ.buffer, biases_qOffset: biasesQ.offset, + w_k: wK.buffer, w_kOffset: wK.offset, + scales_k: scalesK.buffer, scales_kOffset: scalesK.offset, + biases_k: biasesK.buffer, biases_kOffset: biasesK.offset, + w_v: wV.buffer, w_vOffset: wV.offset, + scales_v: scalesV.buffer, scales_vOffset: scalesV.offset, + biases_v: biasesV.buffer, biases_vOffset: biasesV.offset, + out: out.buffer, outOffset: out.offset, + out_q: UInt32(outQ), out_k: UInt32(outK), out_v: UInt32(outV), + in_dim: UInt32(inDim), group_size: UInt32(groupSize), gridSize: grid, threadgroupSize: tg, on: cmd) default: - fatalError("Ops.conv1dCausalStep: unsupported dtype \(x.dtype)") + fatalError("Ops.batchedQkvQgemvInt4Fast: unsupported dtype \(x.dtype)") } } - /// Mamba 2 selective-scan single-token decode step. Updates the - /// per-layer recurrent state `h` in place and writes the output - /// channel vector `y`. `h` lives in fp32 (state accumulates over - /// many decode steps; bf16's 7-bit mantissa drifts fast). One - /// thread per `(head, channel)` — total `nHeads * headDim` threads. + /// M>1 (batched-prefill) sibling of `batchedQkvQgemvInt4Fast`. Reads + /// `x = [M, in_dim]` once into TG memory and produces all 3 + /// projections into THREE separate contiguous output tensors + /// (`qBuf [M, out_q]`, `kBuf [M, out_k]`, `vBuf [M, out_v]`). The + /// split-output layout lets downstream Q/K/V code read each + /// projection as `[M, dim]` without strided views (vs. a single + /// concat buffer). One dispatch replaces three independent qmm + /// dispatches; the `xNorm` DRAM roundtrip is paid once instead of + /// three times. /// - /// See `SSMStateCache` for the storage class that wraps the per-layer - /// `h` buffer; Mamba 2 family files call this through - /// that cache. - public static func ssmStep( - x: Tensor, a: Tensor, b: Tensor, c: Tensor, dt: Tensor, - state h: Tensor, into y: Tensor, - nHeads: Int, headDim: Int, stateDim: Int, - on cmd: MTLCommandBuffer + /// Constraints (mirror `ffai_batched_qkv_qmm_fast`): + /// - `in_dim` MUST be a multiple of 512. + /// - Each of `out_q`, `out_k`, `out_v` MUST be a multiple of 8. + /// - `group_size` MUST be 64. TPG = 64. + /// - All three outputs MUST be pre-zeroed (kernel only writes + /// valid `row0 < out_*` tiles per matrix branch). The wrapper + /// handles this — callers do not need to zero ahead of time. + public static func batchedQkvQmmFast( + x: Tensor, // [M, in_dim] + wQ: Tensor, scalesQ: Tensor, biasesQ: Tensor, + wK: Tensor, scalesK: Tensor, biasesK: Tensor, + wV: Tensor, scalesV: Tensor, biasesV: Tensor, + m: Int, outQ: Int, outK: Int, outV: Int, + on cmd: MTLCommandBuffer, + qBuf: Tensor, // [M, out_q] + kBuf: Tensor, // [M, out_k] + vBuf: Tensor // [M, out_v] ) { - precondition(h.dtype == .f32, "Ops.ssmStep: state h must be f32") - precondition(x.dtype == y.dtype, "Ops.ssmStep: x and y dtype must match") precondition( - a.dtype == x.dtype && b.dtype == x.dtype - && c.dtype == x.dtype && dt.dtype == x.dtype, - "Ops.ssmStep: a/b/c/dt dtype must match x") - let grid = MTLSize(width: nHeads * headDim, height: 1, depth: 1) - let tg = MTLSize(width: 1, height: 1, depth: 1) + wQ.dtype == .u32 && wK.dtype == .u32 && wV.dtype == .u32, + "Ops.batchedQkvQmmFast: w_* must be u32-packed") + let packedPerRow = wQ.shape[1] + let inDim = packedPerRow * 8 + // K and V share the same in_dim as Q — the kernel reads ONE x row + // and walks all three weight tensors at that packed width. A + // mis-packed wK/wV would silently miscompute the K/V projections. + precondition( + wK.shape[1] == packedPerRow && wV.shape[1] == packedPerRow, + "Ops.batchedQkvQmmFast: wK/wV packed width (\(wK.shape[1])/\(wV.shape[1])) must equal wQ (\(packedPerRow))" + ) + precondition( + x.elementCount == m * inDim, + "Ops.batchedQkvQmmFast: x.elementCount \(x.elementCount) ≠ M·inDim \(m * inDim)") + precondition( + qBuf.elementCount == m * outQ, + "Ops.batchedQkvQmmFast: qBuf.elementCount must be M·out_q") + precondition( + kBuf.elementCount == m * outK, + "Ops.batchedQkvQmmFast: kBuf.elementCount must be M·out_k") + precondition( + vBuf.elementCount == m * outV, + "Ops.batchedQkvQmmFast: vBuf.elementCount must be M·out_v") + precondition( + inDim % 512 == 0, + "Ops.batchedQkvQmmFast: in_dim must be a multiple of 512") + precondition( + outQ % 8 == 0 && outK % 8 == 0 && outV % 8 == 0, + "Ops.batchedQkvQmmFast: out_q/k/v must each be a multiple of 8") + // Pre-zero outputs: the kernel tile mechanic depends on the + // buffer starting at zero where the smaller two matrices' tiles + // past their `out_*` count no-op the store. + qBuf.zero() + kBuf.zero() + vBuf.zero() + let groupSize = 64 + let tpg = 64 + let maxOut = max(outQ, max(outK, outV)) + let nTiles = maxOut / 8 + let grid = MTLSize(width: nTiles * tpg, height: m, depth: 3) + let tg = MTLSize(width: tpg, height: 1, depth: 1) switch x.dtype { case .f32: - MetalTileKernels.ssm_step_f32( + MetalTileKernels.ffai_batched_qkv_qmm_fast_f32( x: x.buffer, xOffset: x.offset, - a: a.buffer, aOffset: a.offset, - b: b.buffer, bOffset: b.offset, - c: c.buffer, cOffset: c.offset, - dt: dt.buffer, dtOffset: dt.offset, - h: h.buffer, hOffset: h.offset, - y: y.buffer, yOffset: y.offset, - head_dim: UInt32(headDim), state_dim: UInt32(stateDim), + w_q: wQ.buffer, w_qOffset: wQ.offset, + scales_q: scalesQ.buffer, scales_qOffset: scalesQ.offset, + biases_q: biasesQ.buffer, biases_qOffset: biasesQ.offset, + w_k: wK.buffer, w_kOffset: wK.offset, + scales_k: scalesK.buffer, scales_kOffset: scalesK.offset, + biases_k: biasesK.buffer, biases_kOffset: biasesK.offset, + w_v: wV.buffer, w_vOffset: wV.offset, + scales_v: scalesV.buffer, scales_vOffset: scalesV.offset, + biases_v: biasesV.buffer, biases_vOffset: biasesV.offset, + q_buf: qBuf.buffer, q_bufOffset: qBuf.offset, + k_buf: kBuf.buffer, k_bufOffset: kBuf.offset, + v_buf: vBuf.buffer, v_bufOffset: vBuf.offset, + out_q: UInt32(outQ), out_k: UInt32(outK), out_v: UInt32(outV), + in_dim: UInt32(inDim), group_size: UInt32(groupSize), gridSize: grid, threadgroupSize: tg, on: cmd) case .f16: - MetalTileKernels.ssm_step_f16( + MetalTileKernels.ffai_batched_qkv_qmm_fast_f16( x: x.buffer, xOffset: x.offset, - a: a.buffer, aOffset: a.offset, - b: b.buffer, bOffset: b.offset, - c: c.buffer, cOffset: c.offset, - dt: dt.buffer, dtOffset: dt.offset, - h: h.buffer, hOffset: h.offset, - y: y.buffer, yOffset: y.offset, - head_dim: UInt32(headDim), state_dim: UInt32(stateDim), + w_q: wQ.buffer, w_qOffset: wQ.offset, + scales_q: scalesQ.buffer, scales_qOffset: scalesQ.offset, + biases_q: biasesQ.buffer, biases_qOffset: biasesQ.offset, + w_k: wK.buffer, w_kOffset: wK.offset, + scales_k: scalesK.buffer, scales_kOffset: scalesK.offset, + biases_k: biasesK.buffer, biases_kOffset: biasesK.offset, + w_v: wV.buffer, w_vOffset: wV.offset, + scales_v: scalesV.buffer, scales_vOffset: scalesV.offset, + biases_v: biasesV.buffer, biases_vOffset: biasesV.offset, + q_buf: qBuf.buffer, q_bufOffset: qBuf.offset, + k_buf: kBuf.buffer, k_bufOffset: kBuf.offset, + v_buf: vBuf.buffer, v_bufOffset: vBuf.offset, + out_q: UInt32(outQ), out_k: UInt32(outK), out_v: UInt32(outV), + in_dim: UInt32(inDim), group_size: UInt32(groupSize), gridSize: grid, threadgroupSize: tg, on: cmd) case .bf16: - MetalTileKernels.ssm_step_bf16( + MetalTileKernels.ffai_batched_qkv_qmm_fast_bf16( x: x.buffer, xOffset: x.offset, - a: a.buffer, aOffset: a.offset, - b: b.buffer, bOffset: b.offset, - c: c.buffer, cOffset: c.offset, - dt: dt.buffer, dtOffset: dt.offset, - h: h.buffer, hOffset: h.offset, - y: y.buffer, yOffset: y.offset, - head_dim: UInt32(headDim), state_dim: UInt32(stateDim), + w_q: wQ.buffer, w_qOffset: wQ.offset, + scales_q: scalesQ.buffer, scales_qOffset: scalesQ.offset, + biases_q: biasesQ.buffer, biases_qOffset: biasesQ.offset, + w_k: wK.buffer, w_kOffset: wK.offset, + scales_k: scalesK.buffer, scales_kOffset: scalesK.offset, + biases_k: biasesK.buffer, biases_kOffset: biasesK.offset, + w_v: wV.buffer, w_vOffset: wV.offset, + scales_v: scalesV.buffer, scales_vOffset: scalesV.offset, + biases_v: biasesV.buffer, biases_vOffset: biasesV.offset, + q_buf: qBuf.buffer, q_bufOffset: qBuf.offset, + k_buf: kBuf.buffer, k_bufOffset: kBuf.offset, + v_buf: vBuf.buffer, v_bufOffset: vBuf.offset, + out_q: UInt32(outQ), out_k: UInt32(outK), out_v: UInt32(outV), + in_dim: UInt32(inDim), group_size: UInt32(groupSize), gridSize: grid, threadgroupSize: tg, on: cmd) default: - fatalError("Ops.ssmStep: unsupported dtype \(x.dtype)") + fatalError("Ops.batchedQkvQmmFast: unsupported dtype \(x.dtype)") } } - /// Gated Delta Net (GDN) recurrent single-token decode step. - /// - /// Computes, per decode step and per value head, the recurrence - /// S_t = g_t · S_{t-1} + β_t · k_t · (v_t − k_tᵀ · S_{t-1})ᵀ - /// then the output `o_t = S_t · q_t`. The per-head state matrix - /// `S [Hv, Dv, Dk]` stays fp32 throughout — the gate + rank-1 - /// update accumulate across many steps and bf16 drifts fast. - /// - /// The underlying kernel is **reduction-mode** with strict - /// dispatch-shape invariants (see - /// `crates/metaltile-std/src/ffai/gated_delta.rs` - /// §"DISPATCH INVARIANTS" and `OpsValidation.validateGatedDeltaStep`): - /// TPG = 32 (one simdgroup), one threadgroup per `(dv, n)` pair, - /// `Dk % 32 == 0`, `Hv % Hk == 0`. The `(Dk, Dv, Hk, Hv)` dimensions - /// are runtime `#[constexpr]` scalars — a single PSO serves every - /// model config, no per-tuple specialization. - /// - /// `q` / `k` are expected pre-normalised (the GDN block applies the - /// rmsNorm + scale before calling this — the standard, non-fused - /// kernel variant). The kernel reads `stateIn` and writes a distinct - /// `stateOut`; callers double-buffer via `GDNStateCache.swap()`. + /// Fused 4-output int4 dequant-GEMV in ONE dispatch via + /// `ffai_batched_4_qgemv_fast`. Extends the QKV (3-output) pattern + /// to 4 projections sharing the same input. Used by the Qwen3.5 + /// GDN mixer where the four input projections (qkv, z, b, a) all + /// read the same `xNorm` and can be fused into a single kernel + /// launch. Replaces the 4-dispatch shared-encoder form + /// (`Ops.dequantGemvInt4Four`) with one launch, paying the input + /// DRAM roundtrip once instead of four times. /// - /// All tensors are f32 — the kernel runs the recurrence state in f32 - /// regardless of activation dtype. This wraps `mt_gated_delta_step`, - /// which performs exactly one recurrence step; multi-step prefill - /// loops this call per token at the caller. - public static func gatedDeltaStep( - q: Tensor, k: Tensor, v: Tensor, g: Tensor, beta: Tensor, - stateIn: Tensor, into y: Tensor, stateOut: Tensor, - numKeyHeads: Int, numValueHeads: Int, - keyHeadDim: Int, valueHeadDim: Int, + /// Constraints (mirror `ffai_batched_qkv_qgemv_fast`): + /// - `in_dim` MUST be a multiple of 512. + /// - Each of `out_a / out_b / out_c / out_d` MUST be a multiple + /// of 8. + /// - `group_size` MUST be 64. + /// - TPG = 64. Grid: `[ceil(max(out_*) / 8) * TPG, 1, 4]`. + public static func batched4QgemvInt4Fast( + input x: Tensor, + wA: Tensor, scalesA: Tensor, biasesA: Tensor, outA: Tensor, + wB: Tensor, scalesB: Tensor, biasesB: Tensor, outB: Tensor, + wC: Tensor, scalesC: Tensor, biasesC: Tensor, outC: Tensor, + wD: Tensor, scalesD: Tensor, biasesD: Tensor, outD: Tensor, + groupSize: Int = 64, on cmd: MTLCommandBuffer ) { - // Kernel-invariant validation — see OpsValidation.swift. A bad + precondition( + wA.dtype == .u32 && wB.dtype == .u32 && wC.dtype == .u32 && wD.dtype == .u32, + "Ops.batched4QgemvInt4Fast: w_* must be u32-packed") + let packedPerRow = wA.shape[1] + let inDim = packedPerRow * 8 + // All four weight tensors share the same in_dim — the kernel reads + // ONE x row and walks all four at that packed width. A mis-packed + // wB/wC/wD would silently miscompute its branch's output. + precondition( + wB.shape[1] == packedPerRow && wC.shape[1] == packedPerRow + && wD.shape[1] == packedPerRow, + "Ops.batched4QgemvInt4Fast: wB/wC/wD packed width must equal wA (\(packedPerRow))" + ) + precondition( + x.elementCount == inDim, + "Ops.batched4QgemvInt4Fast: x.elementCount \(x.elementCount) ≠ inDim \(inDim)") + let outA_ = outA.elementCount + let outB_ = outB.elementCount + let outC_ = outC.elementCount + let outD_ = outD.elementCount + precondition( + inDim % 512 == 0, + "Ops.batched4QgemvInt4Fast: in_dim must be a multiple of 512") + precondition( + outA_ % 8 == 0 && outB_ % 8 == 0 && outC_ % 8 == 0 && outD_ % 8 == 0, + "Ops.batched4QgemvInt4Fast: each out_* must be a multiple of 8") + let tpg = 64 + let maxOut = max(max(outA_, outB_), max(outC_, outD_)) + let nTiles = maxOut / 8 + let grid = MTLSize(width: nTiles * tpg, height: 1, depth: 4) + let tg = MTLSize(width: tpg, height: 1, depth: 1) + switch x.dtype { + case .f32: + MetalTileKernels.ffai_batched_4_qgemv_fast_f32( + x: x.buffer, xOffset: x.offset, + w_a: wA.buffer, w_aOffset: wA.offset, + scales_a: scalesA.buffer, scales_aOffset: scalesA.offset, + biases_a: biasesA.buffer, biases_aOffset: biasesA.offset, + w_b: wB.buffer, w_bOffset: wB.offset, + scales_b: scalesB.buffer, scales_bOffset: scalesB.offset, + biases_b: biasesB.buffer, biases_bOffset: biasesB.offset, + w_c: wC.buffer, w_cOffset: wC.offset, + scales_c: scalesC.buffer, scales_cOffset: scalesC.offset, + biases_c: biasesC.buffer, biases_cOffset: biasesC.offset, + w_d: wD.buffer, w_dOffset: wD.offset, + scales_d: scalesD.buffer, scales_dOffset: scalesD.offset, + biases_d: biasesD.buffer, biases_dOffset: biasesD.offset, + a_out: outA.buffer, a_outOffset: outA.offset, + b_out: outB.buffer, b_outOffset: outB.offset, + c_out: outC.buffer, c_outOffset: outC.offset, + d_out: outD.buffer, d_outOffset: outD.offset, + out_a: UInt32(outA_), out_b: UInt32(outB_), + out_c: UInt32(outC_), out_d: UInt32(outD_), + in_dim: UInt32(inDim), group_size: UInt32(groupSize), + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.ffai_batched_4_qgemv_fast_f16( + x: x.buffer, xOffset: x.offset, + w_a: wA.buffer, w_aOffset: wA.offset, + scales_a: scalesA.buffer, scales_aOffset: scalesA.offset, + biases_a: biasesA.buffer, biases_aOffset: biasesA.offset, + w_b: wB.buffer, w_bOffset: wB.offset, + scales_b: scalesB.buffer, scales_bOffset: scalesB.offset, + biases_b: biasesB.buffer, biases_bOffset: biasesB.offset, + w_c: wC.buffer, w_cOffset: wC.offset, + scales_c: scalesC.buffer, scales_cOffset: scalesC.offset, + biases_c: biasesC.buffer, biases_cOffset: biasesC.offset, + w_d: wD.buffer, w_dOffset: wD.offset, + scales_d: scalesD.buffer, scales_dOffset: scalesD.offset, + biases_d: biasesD.buffer, biases_dOffset: biasesD.offset, + a_out: outA.buffer, a_outOffset: outA.offset, + b_out: outB.buffer, b_outOffset: outB.offset, + c_out: outC.buffer, c_outOffset: outC.offset, + d_out: outD.buffer, d_outOffset: outD.offset, + out_a: UInt32(outA_), out_b: UInt32(outB_), + out_c: UInt32(outC_), out_d: UInt32(outD_), + in_dim: UInt32(inDim), group_size: UInt32(groupSize), + gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.ffai_batched_4_qgemv_fast_bf16( + x: x.buffer, xOffset: x.offset, + w_a: wA.buffer, w_aOffset: wA.offset, + scales_a: scalesA.buffer, scales_aOffset: scalesA.offset, + biases_a: biasesA.buffer, biases_aOffset: biasesA.offset, + w_b: wB.buffer, w_bOffset: wB.offset, + scales_b: scalesB.buffer, scales_bOffset: scalesB.offset, + biases_b: biasesB.buffer, biases_bOffset: biasesB.offset, + w_c: wC.buffer, w_cOffset: wC.offset, + scales_c: scalesC.buffer, scales_cOffset: scalesC.offset, + biases_c: biasesC.buffer, biases_cOffset: biasesC.offset, + w_d: wD.buffer, w_dOffset: wD.offset, + scales_d: scalesD.buffer, scales_dOffset: scalesD.offset, + biases_d: biasesD.buffer, biases_dOffset: biasesD.offset, + a_out: outA.buffer, a_outOffset: outA.offset, + b_out: outB.buffer, b_outOffset: outB.offset, + c_out: outC.buffer, c_outOffset: outC.offset, + d_out: outD.buffer, d_outOffset: outD.offset, + out_a: UInt32(outA_), out_b: UInt32(outB_), + out_c: UInt32(outC_), out_d: UInt32(outD_), + in_dim: UInt32(inDim), group_size: UInt32(groupSize), + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError("Ops.batched4QgemvInt4Fast: unsupported dtype \(x.dtype)") + } + } + + /// M>1 (batched-prefill) sibling of `batched4QgemvInt4Fast`. Reads + /// `x = [M, in_dim]` once into TG memory and produces all 4 + /// outputs (each `[M, out_*]`) in a single dispatch. Eliminates + /// the 3 redundant input DRAM reads of the unfused `callMany` + /// chain used by the GDN mixer's batched forward (qkv, z, b, a all + /// project from the same `xNorm`). + /// + /// Output tensors are caller-allocated and indexed independently + /// (`outA / outB / outC / outD`, each `[M, out_*]`). + /// + /// Constraints (mirror `ffai_batched_4_qgemv_fast`): + /// - `in_dim` MUST be a multiple of 512. + /// - Each of `out_a / out_b / out_c / out_d` MUST be a multiple + /// of 8. + /// - `group_size` MUST be 64. TPG = 64. + public static func batched4QmmFast( + input x: Tensor, m: Int, + wA: Tensor, scalesA: Tensor, biasesA: Tensor, outA: Tensor, + wB: Tensor, scalesB: Tensor, biasesB: Tensor, outB: Tensor, + wC: Tensor, scalesC: Tensor, biasesC: Tensor, outC: Tensor, + wD: Tensor, scalesD: Tensor, biasesD: Tensor, outD: Tensor, + groupSize: Int = 64, + on cmd: MTLCommandBuffer + ) { + precondition( + wA.dtype == .u32 && wB.dtype == .u32 && wC.dtype == .u32 && wD.dtype == .u32, + "Ops.batched4QmmFast: w_* must be u32-packed") + let packedPerRow = wA.shape[1] + let inDim = packedPerRow * 8 + // All four weight tensors share the same in_dim — the kernel reads + // ONE x row and walks all four at that packed width. A mis-packed + // wB/wC/wD would silently miscompute its branch's output. + precondition( + wB.shape[1] == packedPerRow && wC.shape[1] == packedPerRow + && wD.shape[1] == packedPerRow, + "Ops.batched4QmmFast: wB/wC/wD packed width must equal wA (\(packedPerRow))" + ) + precondition( + x.elementCount == m * inDim, + "Ops.batched4QmmFast: x size \(x.elementCount) ≠ M·inDim \(m * inDim)") + let outADim = wA.shape[0] + let outBDim = wB.shape[0] + let outCDim = wC.shape[0] + let outDDim = wD.shape[0] + precondition( + outA.elementCount == m * outADim, + "Ops.batched4QmmFast: outA size \(outA.elementCount) ≠ M·outADim \(m * outADim)") + precondition( + outB.elementCount == m * outBDim, + "Ops.batched4QmmFast: outB size \(outB.elementCount) ≠ M·outBDim \(m * outBDim)") + precondition( + outC.elementCount == m * outCDim, + "Ops.batched4QmmFast: outC size \(outC.elementCount) ≠ M·outCDim \(m * outCDim)") + precondition( + outD.elementCount == m * outDDim, + "Ops.batched4QmmFast: outD size \(outD.elementCount) ≠ M·outDDim \(m * outDDim)") + precondition( + inDim % 512 == 0, + "Ops.batched4QmmFast: in_dim must be a multiple of 512") + precondition( + outADim % 8 == 0 && outBDim % 8 == 0 && outCDim % 8 == 0 && outDDim % 8 == 0, + "Ops.batched4QmmFast: each out_* must be a multiple of 8") + let tpg = 64 + let maxOut = max(max(outADim, outBDim), max(outCDim, outDDim)) + let nTiles = maxOut / 8 + let grid = MTLSize(width: nTiles * tpg, height: m, depth: 4) + let tg = MTLSize(width: tpg, height: 1, depth: 1) + switch x.dtype { + case .f32: + MetalTileKernels.ffai_batched_4_qmm_fast_f32( + x: x.buffer, xOffset: x.offset, + w_a: wA.buffer, w_aOffset: wA.offset, + scales_a: scalesA.buffer, scales_aOffset: scalesA.offset, + biases_a: biasesA.buffer, biases_aOffset: biasesA.offset, + w_b: wB.buffer, w_bOffset: wB.offset, + scales_b: scalesB.buffer, scales_bOffset: scalesB.offset, + biases_b: biasesB.buffer, biases_bOffset: biasesB.offset, + w_c: wC.buffer, w_cOffset: wC.offset, + scales_c: scalesC.buffer, scales_cOffset: scalesC.offset, + biases_c: biasesC.buffer, biases_cOffset: biasesC.offset, + w_d: wD.buffer, w_dOffset: wD.offset, + scales_d: scalesD.buffer, scales_dOffset: scalesD.offset, + biases_d: biasesD.buffer, biases_dOffset: biasesD.offset, + a_buf: outA.buffer, a_bufOffset: outA.offset, + b_buf: outB.buffer, b_bufOffset: outB.offset, + c_buf: outC.buffer, c_bufOffset: outC.offset, + d_buf: outD.buffer, d_bufOffset: outD.offset, + out_a: UInt32(outADim), out_b: UInt32(outBDim), + out_c: UInt32(outCDim), out_d: UInt32(outDDim), + in_dim: UInt32(inDim), group_size: UInt32(groupSize), + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.ffai_batched_4_qmm_fast_f16( + x: x.buffer, xOffset: x.offset, + w_a: wA.buffer, w_aOffset: wA.offset, + scales_a: scalesA.buffer, scales_aOffset: scalesA.offset, + biases_a: biasesA.buffer, biases_aOffset: biasesA.offset, + w_b: wB.buffer, w_bOffset: wB.offset, + scales_b: scalesB.buffer, scales_bOffset: scalesB.offset, + biases_b: biasesB.buffer, biases_bOffset: biasesB.offset, + w_c: wC.buffer, w_cOffset: wC.offset, + scales_c: scalesC.buffer, scales_cOffset: scalesC.offset, + biases_c: biasesC.buffer, biases_cOffset: biasesC.offset, + w_d: wD.buffer, w_dOffset: wD.offset, + scales_d: scalesD.buffer, scales_dOffset: scalesD.offset, + biases_d: biasesD.buffer, biases_dOffset: biasesD.offset, + a_buf: outA.buffer, a_bufOffset: outA.offset, + b_buf: outB.buffer, b_bufOffset: outB.offset, + c_buf: outC.buffer, c_bufOffset: outC.offset, + d_buf: outD.buffer, d_bufOffset: outD.offset, + out_a: UInt32(outADim), out_b: UInt32(outBDim), + out_c: UInt32(outCDim), out_d: UInt32(outDDim), + in_dim: UInt32(inDim), group_size: UInt32(groupSize), + gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.ffai_batched_4_qmm_fast_bf16( + x: x.buffer, xOffset: x.offset, + w_a: wA.buffer, w_aOffset: wA.offset, + scales_a: scalesA.buffer, scales_aOffset: scalesA.offset, + biases_a: biasesA.buffer, biases_aOffset: biasesA.offset, + w_b: wB.buffer, w_bOffset: wB.offset, + scales_b: scalesB.buffer, scales_bOffset: scalesB.offset, + biases_b: biasesB.buffer, biases_bOffset: biasesB.offset, + w_c: wC.buffer, w_cOffset: wC.offset, + scales_c: scalesC.buffer, scales_cOffset: scalesC.offset, + biases_c: biasesC.buffer, biases_cOffset: biasesC.offset, + w_d: wD.buffer, w_dOffset: wD.offset, + scales_d: scalesD.buffer, scales_dOffset: scalesD.offset, + biases_d: biasesD.buffer, biases_dOffset: biasesD.offset, + a_buf: outA.buffer, a_bufOffset: outA.offset, + b_buf: outB.buffer, b_bufOffset: outB.offset, + c_buf: outC.buffer, c_bufOffset: outC.offset, + d_buf: outD.buffer, d_bufOffset: outD.offset, + out_a: UInt32(outADim), out_b: UInt32(outBDim), + out_c: UInt32(outCDim), out_d: UInt32(outDDim), + in_dim: UInt32(inDim), group_size: UInt32(groupSize), + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError("Ops.batched4QmmFast: unsupported dtype \(x.dtype)") + } + } + + /// Per-expert indexed int4 dequant-GEMV. The caller stacks every + /// expert's weight slab into one `[nExperts, outDim, inDim/8]` + /// u32-packed tensor (and matching scales / biases stacks); the + /// kernel reads `expertIndex[0]` to pick the slab to dequantize on + /// this call. Paired with `Ops.moeRouterTopK`, the top-K indices + /// stay GPU-resident and the host-side `route → dispatch expert + /// gemv` sync at every MoE layer collapses to one + /// `expertIndex.buffer + offset = slot · 4` view per slot. + public static func dequantGemvInt4ExpertIndexed( + weightsStacked: Tensor, scalesStacked: Tensor, biasesStacked: Tensor, + input: Tensor, expertIndex: Tensor, + groupSize: Int = 64, + on cmd: MTLCommandBuffer, + into out: Tensor + ) { + precondition( + weightsStacked.shape.count == 3, + "Ops.dequantGemvInt4ExpertIndexed: weightsStacked must be [nExperts, outDim, inDim/8]") + precondition( + weightsStacked.dtype == .u32, + "Ops.dequantGemvInt4ExpertIndexed: weightsStacked must be u32 (packed)") + precondition( + expertIndex.dtype == .u32 && expertIndex.elementCount == 1, + "Ops.dequantGemvInt4ExpertIndexed: expertIndex must be [1] u32") + precondition( + scalesStacked.dtype == input.dtype && biasesStacked.dtype == input.dtype, + "Ops.dequantGemvInt4ExpertIndexed: scales / biases dtype must match input") + precondition( + out.dtype == input.dtype, + "Ops.dequantGemvInt4ExpertIndexed: out dtype must match input") + let outDim = weightsStacked.shape[1] + let packedPerRow = weightsStacked.shape[2] + let inDim = packedPerRow * 8 // int4 packs 8 weights per u32 word + precondition( + input.elementCount == inDim, + "Ops.dequantGemvInt4ExpertIndexed: input \(input.elementCount) ≠ inDim \(inDim)") + precondition( + out.elementCount == outDim, + "Ops.dequantGemvInt4ExpertIndexed: out \(out.elementCount) ≠ outDim \(outDim)") + // Kernel runs in Reduction mode with tpg=32 (one simdgroup per + // output row). `dispatchThreads` counts THREADS not threadgroups + // — total threads = outDim · 32 gives outDim threadgroups, one + // per output row. + let tg = MTLSize(width: 32, height: 1, depth: 1) + let grid = MTLSize(width: outDim * 32, height: 1, depth: 1) + let inDimU = UInt32(inDim) + let outDimU = UInt32(outDim) + let groupSizeU = UInt32(groupSize) + switch input.dtype { + case .f32: + MetalTileKernels.dequant_gemv_int4_expert_indexed_f32( + weights_stacked: weightsStacked.buffer, weights_stackedOffset: weightsStacked.offset, + scales_stacked: scalesStacked.buffer, scales_stackedOffset: scalesStacked.offset, + biases_stacked: biasesStacked.buffer, biases_stackedOffset: biasesStacked.offset, + input: input.buffer, inputOffset: input.offset, + expert_index: expertIndex.buffer, expert_indexOffset: expertIndex.offset, + output: out.buffer, outputOffset: out.offset, + in_dim: inDimU, out_dim: outDimU, group_size: groupSizeU, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.dequant_gemv_int4_expert_indexed_f16( + weights_stacked: weightsStacked.buffer, weights_stackedOffset: weightsStacked.offset, + scales_stacked: scalesStacked.buffer, scales_stackedOffset: scalesStacked.offset, + biases_stacked: biasesStacked.buffer, biases_stackedOffset: biasesStacked.offset, + input: input.buffer, inputOffset: input.offset, + expert_index: expertIndex.buffer, expert_indexOffset: expertIndex.offset, + output: out.buffer, outputOffset: out.offset, + in_dim: inDimU, out_dim: outDimU, group_size: groupSizeU, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.dequant_gemv_int4_expert_indexed_bf16( + weights_stacked: weightsStacked.buffer, weights_stackedOffset: weightsStacked.offset, + scales_stacked: scalesStacked.buffer, scales_stackedOffset: scalesStacked.offset, + biases_stacked: biasesStacked.buffer, biases_stackedOffset: biasesStacked.offset, + input: input.buffer, inputOffset: input.offset, + expert_index: expertIndex.buffer, expert_indexOffset: expertIndex.offset, + output: out.buffer, outputOffset: out.offset, + in_dim: inDimU, out_dim: outDimU, group_size: groupSizeU, + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError( + "Ops.dequantGemvInt4ExpertIndexed: unsupported input dtype \(input.dtype)") + } + } + + /// Encoder-batched form of `dequantGemvInt4ExpertIndexed`. All N + /// calls must share `(inDim, outDim, groupSize, dtype)` so they + /// reuse the same PSO + constexpr binding; only the per-call + /// `(weights, scales, biases, input, expertIndex, output)` rotate + /// per dispatch. Saves N-1 encoder begin/end pairs versus N + /// independent calls. At Qwen3.6-A3B's topK=8 the gate+up phase + /// becomes 16 calls on one encoder and the down phase 8 calls on + /// a second encoder per MoE layer. + public static func dequantGemvInt4ExpertIndexedMany( + weightsStacked: [Tensor], scalesStacked: [Tensor], biasesStacked: [Tensor], + inputs: [Tensor], expertIndices: [Tensor], outputs: [Tensor], + groupSize: Int = 64, + on cmd: MTLCommandBuffer + ) { + let n = weightsStacked.count + precondition( + scalesStacked.count == n && biasesStacked.count == n + && inputs.count == n && expertIndices.count == n + && outputs.count == n, + "Ops.dequantGemvInt4ExpertIndexedMany: count mismatch") + guard n > 0 else { return } + let dtype = inputs[0].dtype + let psoName: String + switch dtype { + case .f32: psoName = "dequant_gemv_int4_expert_indexed_f32" + case .f16: psoName = "dequant_gemv_int4_expert_indexed_f16" + case .bf16: psoName = "dequant_gemv_int4_expert_indexed_bf16" + default: + fatalError( + "Ops.dequantGemvInt4ExpertIndexedMany: unsupported dtype \(dtype)") + } + let outDim = weightsStacked[0].shape[1] + let packedPerRow = weightsStacked[0].shape[2] + let inDim = packedPerRow * 8 + var inDimV = UInt32(inDim) + var outDimV = UInt32(outDim) + var groupSizeV = UInt32(groupSize) + let pso = PSOCache.shared.pipelineState(for: psoName) + guard let enc = cmd.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(pso) + // Shared constexprs (kernel buffer indices 6, 7, 8) set once. + enc.setBytes(&inDimV, length: 4, index: 6) + enc.setBytes(&outDimV, length: 4, index: 7) + enc.setBytes(&groupSizeV, length: 4, index: 8) + let tg = MTLSize(width: 32, height: 1, depth: 1) + let grid = MTLSize(width: outDim * 32, height: 1, depth: 1) + for i in 0 ..< n { + precondition( + weightsStacked[i].shape[1] == outDim + && weightsStacked[i].shape[2] == packedPerRow, + "Ops.dequantGemvInt4ExpertIndexedMany: weight shape varies at \(i)") + precondition( + inputs[i].dtype == dtype && outputs[i].dtype == dtype, + "Ops.dequantGemvInt4ExpertIndexedMany: dtype varies at \(i)") + precondition( + expertIndices[i].dtype == .u32 && expertIndices[i].elementCount == 1, + "Ops.dequantGemvInt4ExpertIndexedMany: expertIndices[\(i)] must be [1] u32") + enc.setBuffer(weightsStacked[i].buffer, offset: weightsStacked[i].offset, index: 0) + enc.setBuffer(scalesStacked[i].buffer, offset: scalesStacked[i].offset, index: 1) + enc.setBuffer(biasesStacked[i].buffer, offset: biasesStacked[i].offset, index: 2) + enc.setBuffer(inputs[i].buffer, offset: inputs[i].offset, index: 3) + enc.setBuffer(expertIndices[i].buffer, offset: expertIndices[i].offset, index: 4) + enc.setBuffer(outputs[i].buffer, offset: outputs[i].offset, index: 5) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + } + enc.endEncoding() + } + + /// MoE top-K router. Reads `[nExperts]` raw logits and writes + /// `[k]` u32 expert indices + `[k]` weights of the same dtype. + /// `normTopkProb` matches Qwen3-MoE convention when `true` + /// (softmax restricted to the chosen k, weights renormalise to + /// sum-to-1); `false` matches the Qwen3-Next style (softmax over + /// all experts, then pick top-k without renormalisation). + /// + /// Single-row form — `logits: [nExperts]`. Internally calls the + /// T-batched variant with `t = 1`. + public static func moeRouterTopK( + logits: Tensor, indicesOut: Tensor, weightsOut: Tensor, + nExperts: Int, k: Int, normTopkProb: Bool, + on cmd: MTLCommandBuffer + ) { + moeRouterTopKMany( + logits: logits, indicesOut: indicesOut, weightsOut: weightsOut, + t: 1, nExperts: nExperts, k: k, normTopkProb: normTopkProb, + on: cmd) + } + + /// T-batched MoE top-K router. `logits: [T, nExperts]`, + /// `indicesOut: [T, k]` u32, `weightsOut: [T, k]` (matching logits + /// dtype). The kernel iterates `T` rows along `program_id<0>()` + /// (one threadgroup per row), so a single dispatch covers all + /// prefill rows without the host commit + wait + CPU + /// `MoERouter.route` round-trip that the per-row form forced. + public static func moeRouterTopKMany( + logits: Tensor, indicesOut: Tensor, weightsOut: Tensor, + t: Int, nExperts: Int, k: Int, normTopkProb: Bool, + on cmd: MTLCommandBuffer + ) { + precondition(t > 0, "Ops.moeRouterTopKMany: t must be positive") + precondition( + logits.elementCount == t * nExperts, + "Ops.moeRouterTopKMany: logits must be [T·nExperts]") + precondition( + indicesOut.elementCount == t * k && indicesOut.dtype == .u32, + "Ops.moeRouterTopKMany: indicesOut must be [T·k] u32") + precondition( + weightsOut.elementCount == t * k && weightsOut.dtype == logits.dtype, + "Ops.moeRouterTopKMany: weightsOut must be [T·k] matching logits dtype") + // Kernel pins tpg = 32 (one simdgroup per token row, Reduction + // mode). Grid total threads = T · 32 → T threadgroups. + let tg = MTLSize(width: 32, height: 1, depth: 1) + let grid = MTLSize(width: t * 32, height: 1, depth: 1) + let normFlag: UInt32 = normTopkProb ? 1 : 0 + switch logits.dtype { + case .f32: + MetalTileKernels.mt_moe_router_topk_f32( + router_logits: logits.buffer, router_logitsOffset: logits.offset, + indices_out: indicesOut.buffer, indices_outOffset: indicesOut.offset, + weights_out: weightsOut.buffer, weights_outOffset: weightsOut.offset, + n_experts: UInt32(nExperts), k: UInt32(k), norm_topk_prob: normFlag, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.mt_moe_router_topk_f16( + router_logits: logits.buffer, router_logitsOffset: logits.offset, + indices_out: indicesOut.buffer, indices_outOffset: indicesOut.offset, + weights_out: weightsOut.buffer, weights_outOffset: weightsOut.offset, + n_experts: UInt32(nExperts), k: UInt32(k), norm_topk_prob: normFlag, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.mt_moe_router_topk_bf16( + router_logits: logits.buffer, router_logitsOffset: logits.offset, + indices_out: indicesOut.buffer, indices_outOffset: indicesOut.offset, + weights_out: weightsOut.buffer, weights_outOffset: weightsOut.offset, + n_experts: UInt32(nExperts), k: UInt32(k), norm_topk_prob: normFlag, + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError( + "Ops.moeRouterTopKMany: unsupported logits dtype \(logits.dtype)") + } + } + + /// GPU argmax over a 1D logits tensor. Caller supplies a 1-element + /// u32 output buffer. Uses the cooperative 256-thread Reduction + /// kernel — one threadgroup, ~80-300 KB / vocab logits in registers. + /// Mamba 2 / Mamba 1D depthwise causal-conv step — streaming + /// decode form. One thread per channel. `state` is the rolling + /// window of the last `kernelSize - 1` inputs (shape + /// `[kernelSize - 1, nChannels]`); shifted in-place after compute. + /// Activation (Mamba 2 follows the conv with SiLU) is the caller's + /// concern — kept separate for composability. + public static func conv1dCausalStep( + x: Tensor, w: Tensor, b: Tensor, + state: Tensor, into y: Tensor, + nChannels: Int, kernelSize: Int, + on cmd: MTLCommandBuffer + ) { + precondition( + w.dtype == x.dtype && b.dtype == x.dtype + && state.dtype == x.dtype && y.dtype == x.dtype, + "Ops.conv1dCausalStep: every tensor must share dtype") + let grid = MTLSize(width: nChannels, height: 1, depth: 1) + let tg = MTLSize(width: 1, height: 1, depth: 1) + switch x.dtype { + case .f32: + MetalTileKernels.conv1d_causal_step_f32( + x: x.buffer, xOffset: x.offset, + w: w.buffer, wOffset: w.offset, + b: b.buffer, bOffset: b.offset, + state: state.buffer, stateOffset: state.offset, + y: y.buffer, yOffset: y.offset, + n_channels: UInt32(nChannels), kernel_size: UInt32(kernelSize), + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.conv1d_causal_step_f16( + x: x.buffer, xOffset: x.offset, + w: w.buffer, wOffset: w.offset, + b: b.buffer, bOffset: b.offset, + state: state.buffer, stateOffset: state.offset, + y: y.buffer, yOffset: y.offset, + n_channels: UInt32(nChannels), kernel_size: UInt32(kernelSize), + gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.conv1d_causal_step_bf16( + x: x.buffer, xOffset: x.offset, + w: w.buffer, wOffset: w.offset, + b: b.buffer, bOffset: b.offset, + state: state.buffer, stateOffset: state.offset, + y: y.buffer, yOffset: y.offset, + n_channels: UInt32(nChannels), kernel_size: UInt32(kernelSize), + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError("Ops.conv1dCausalStep: unsupported dtype \(x.dtype)") + } + } + + /// Batched depthwise causal conv1d + SiLU + cast-to-f32. Sweeps + /// `t` tokens in ONE dispatch with the conv state held in + /// per-channel registers across the sweep — `stateIn` is read + /// once at the top, `stateOut` is written once at the bottom. + /// Replaces the per-token `conv1dCausalStep` + `siluCastToF32` + /// T-loop in `Qwen35GDNMixer.forwardManyChunked`. + /// + /// `convKernel` must be 4 (Qwen3.5 / 3.6, Mamba 2, NemotronH all + /// use kernel=4; the kernel signature carries it as a constexpr + /// for shape uniformity but the body is K=4 hardcoded). + /// + /// State in-place safety: each grid thread owns one channel; + /// reads `stateIn[*, c]` once at the start of the sweep and + /// writes `stateOut[*, c]` once at the end. Passing the same + /// buffer for both is safe. + public static func conv1dCausalStepSiluCastMany( + src: Tensor, w: Tensor, b: Tensor, + stateIn: Tensor, outF32: Tensor, stateOut: Tensor, + t: Int, convDim: Int, convKernel: Int, + on cmd: MTLCommandBuffer + ) { + precondition( + convKernel == 4, + "Ops.conv1dCausalStepSiluCastMany: requires conv_kernel = 4") + precondition( + src.dtype == w.dtype && src.dtype == b.dtype + && src.dtype == stateIn.dtype && src.dtype == stateOut.dtype, + "Ops.conv1dCausalStepSiluCastMany: dtype mismatch") + precondition( + outF32.dtype == .f32, + "Ops.conv1dCausalStepSiluCastMany: outF32 must be .f32") + let grid = MTLSize(width: convDim, height: 1, depth: 1) + let tg = MTLSize(width: min(convDim, 256), height: 1, depth: 1) + let tLenU = UInt32(t) + let cdU = UInt32(convDim) + let ckU = UInt32(convKernel) + switch src.dtype { + case .f32: + MetalTileKernels.ffai_conv1d_causal_step_silu_cast_many_f32( + src: src.buffer, srcOffset: src.offset, + w: w.buffer, wOffset: w.offset, + b: b.buffer, bOffset: b.offset, + state_in: stateIn.buffer, state_inOffset: stateIn.offset, + out_f32: outF32.buffer, out_f32Offset: outF32.offset, + state_out: stateOut.buffer, state_outOffset: stateOut.offset, + t_len: tLenU, conv_dim: cdU, conv_kernel: ckU, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.ffai_conv1d_causal_step_silu_cast_many_f16( + src: src.buffer, srcOffset: src.offset, + w: w.buffer, wOffset: w.offset, + b: b.buffer, bOffset: b.offset, + state_in: stateIn.buffer, state_inOffset: stateIn.offset, + out_f32: outF32.buffer, out_f32Offset: outF32.offset, + state_out: stateOut.buffer, state_outOffset: stateOut.offset, + t_len: tLenU, conv_dim: cdU, conv_kernel: ckU, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.ffai_conv1d_causal_step_silu_cast_many_bf16( + src: src.buffer, srcOffset: src.offset, + w: w.buffer, wOffset: w.offset, + b: b.buffer, bOffset: b.offset, + state_in: stateIn.buffer, state_inOffset: stateIn.offset, + out_f32: outF32.buffer, out_f32Offset: outF32.offset, + state_out: stateOut.buffer, state_outOffset: stateOut.offset, + t_len: tLenU, conv_dim: cdU, conv_kernel: ckU, + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError( + "Ops.conv1dCausalStepSiluCastMany: unsupported dtype \(src.dtype)") + } + } + + /// Mamba 2 selective-scan single-token decode step. Updates the + /// per-layer recurrent state `h` in place and writes the output + /// channel vector `y`. `h` lives in fp32 (state accumulates over + /// many decode steps; bf16's 7-bit mantissa drifts fast). One + /// thread per `(head, channel)` — total `nHeads * headDim` threads. + /// + /// See `SSMStateCache` for the storage class that wraps the per-layer + /// `h` buffer; Mamba 2 family files call this through + /// that cache. + public static func ssmStep( + x: Tensor, a: Tensor, b: Tensor, c: Tensor, dt: Tensor, + state h: Tensor, into y: Tensor, + nHeads: Int, headDim: Int, stateDim: Int, + on cmd: MTLCommandBuffer + ) { + precondition(h.dtype == .f32, "Ops.ssmStep: state h must be f32") + precondition(x.dtype == y.dtype, "Ops.ssmStep: x and y dtype must match") + precondition( + a.dtype == x.dtype && b.dtype == x.dtype + && c.dtype == x.dtype && dt.dtype == x.dtype, + "Ops.ssmStep: a/b/c/dt dtype must match x") + let grid = MTLSize(width: nHeads * headDim, height: 1, depth: 1) + let tg = MTLSize(width: 1, height: 1, depth: 1) + switch x.dtype { + case .f32: + MetalTileKernels.ssm_step_f32( + x: x.buffer, xOffset: x.offset, + a: a.buffer, aOffset: a.offset, + b: b.buffer, bOffset: b.offset, + c: c.buffer, cOffset: c.offset, + dt: dt.buffer, dtOffset: dt.offset, + h: h.buffer, hOffset: h.offset, + y: y.buffer, yOffset: y.offset, + head_dim: UInt32(headDim), state_dim: UInt32(stateDim), + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.ssm_step_f16( + x: x.buffer, xOffset: x.offset, + a: a.buffer, aOffset: a.offset, + b: b.buffer, bOffset: b.offset, + c: c.buffer, cOffset: c.offset, + dt: dt.buffer, dtOffset: dt.offset, + h: h.buffer, hOffset: h.offset, + y: y.buffer, yOffset: y.offset, + head_dim: UInt32(headDim), state_dim: UInt32(stateDim), + gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.ssm_step_bf16( + x: x.buffer, xOffset: x.offset, + a: a.buffer, aOffset: a.offset, + b: b.buffer, bOffset: b.offset, + c: c.buffer, cOffset: c.offset, + dt: dt.buffer, dtOffset: dt.offset, + h: h.buffer, hOffset: h.offset, + y: y.buffer, yOffset: y.offset, + head_dim: UInt32(headDim), state_dim: UInt32(stateDim), + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError("Ops.ssmStep: unsupported dtype \(x.dtype)") + } + } + + /// Gated Delta Net (GDN) recurrent single-token decode step. + /// + /// Computes, per decode step and per value head, the recurrence + /// S_t = g_t · S_{t-1} + β_t · k_t · (v_t − k_tᵀ · S_{t-1})ᵀ + /// then the output `o_t = S_t · q_t`. The per-head state matrix + /// `S [Hv, Dv, Dk]` stays fp32 throughout — the gate + rank-1 + /// update accumulate across many steps and bf16 drifts fast. + /// + /// The underlying kernel is **reduction-mode** with strict + /// dispatch-shape invariants (see + /// `crates/metaltile-std/src/ffai/gated_delta.rs` + /// §"DISPATCH INVARIANTS" and `OpsValidation.validateGatedDeltaStep`): + /// TPG = 32 (one simdgroup), one threadgroup per `(dv, n)` pair, + /// `Dk % 32 == 0`, `Hv % Hk == 0`. The `(Dk, Dv, Hk, Hv)` dimensions + /// are runtime `#[constexpr]` scalars — a single PSO serves every + /// model config, no per-tuple specialization. + /// + /// `q` / `k` are expected pre-normalised (the GDN block applies the + /// rmsNorm + scale before calling this — the standard, non-fused + /// kernel variant). The kernel reads `stateIn` and writes a distinct + /// `stateOut`; callers double-buffer via `GDNStateCache.swap()`. + /// + /// All tensors are f32 — the kernel runs the recurrence state in f32 + /// regardless of activation dtype. This wraps `mt_gated_delta_step`, + /// which performs exactly one recurrence step; multi-step prefill + /// loops this call per token at the caller. + public static func gatedDeltaStep( + q: Tensor, k: Tensor, v: Tensor, g: Tensor, beta: Tensor, + stateIn: Tensor, into y: Tensor, stateOut: Tensor, + numKeyHeads: Int, numValueHeads: Int, + keyHeadDim: Int, valueHeadDim: Int, + on cmd: MTLCommandBuffer + ) { + // Kernel-invariant validation — see OpsValidation.swift. A bad // dispatch shape on a reduction kernel ranges from silent // miscompute to a non-preemptive GPU pin. if let reason = OpsValidation.validateGatedDeltaStep( @@ -1929,6 +3289,113 @@ public enum Ops { } } + /// Append the current token's K AND V rows to their caches in ONE + /// compute encoder. The cache buffers, layout, and writing thread + /// are identical to `kvCacheUpdate` — this just amortises the + /// encoder begin/end across both projections (saving one pair per + /// attention layer per decode token). + public static func kvCacheUpdateKV( + kSrc: Tensor, kCache: Tensor, + vSrc: Tensor, vCache: Tensor, + nKVHeads: Int, headDim: Int, maxSeq: Int, position: Int, + on cmd: MTLCommandBuffer + ) { + precondition( + kSrc.dtype == kCache.dtype && vSrc.dtype == vCache.dtype + && kSrc.dtype == vSrc.dtype, + "Ops.kvCacheUpdateKV: dtype mismatch") + let total = nKVHeads * headDim + let (grid, tg) = elementwiseGrid(total) + let hd = UInt32(headDim) + let ms = UInt32(maxSeq) + let pos = UInt32(position) + let psoName: String + switch kSrc.dtype { + case .f32: psoName = "kv_cache_update_f32" + case .f16: psoName = "kv_cache_update_f16" + case .bf16: psoName = "kv_cache_update_bf16" + default: fatalError("Ops.kvCacheUpdateKV: unsupported dtype \(kSrc.dtype)") + } + let pso = PSOCache.shared.pipelineState(for: psoName) + guard let enc = cmd.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(pso) + @inline(__always) + func dispatch(_ src: Tensor, _ cache: Tensor) { + enc.setBuffer(src.buffer, offset: src.offset, index: 0) + enc.setBuffer(cache.buffer, offset: cache.offset, index: 1) + var headDimV = hd + var maxSeqV = ms + var positionV = pos + enc.setBytes(&headDimV, length: 4, index: 2) + enc.setBytes(&maxSeqV, length: 4, index: 3) + enc.setBytes(&positionV, length: 4, index: 4) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + } + dispatch(kSrc, kCache) + dispatch(vSrc, vCache) + enc.endEncoding() + } + + /// Batched K + V cache append: writes T tokens' K rows AND V rows + /// into their respective caches in ONE shared encoder (2 batched + /// dispatches). Each per-token write goes to a slot indexed by + /// `positions[t]`. Replaces the per-token `kvCacheUpdateKV` loop + /// at `decodeMany` / `forwardMany` prefill — at T=512 × 10 attn + /// layers that's ≈ 5100 fewer dispatches per prefill call. + /// + /// Shapes: + /// - `kSrc` / `vSrc`: `[T, nKVHeads, headDim]` (flat). + /// - `kCache` / `vCache`: `[nKVHeads, maxSeq, headDim]`. + /// - `positions`: `[T]` u32 slot indices into the caches. + public static func kvCacheUpdateKVMany( + kSrc: Tensor, kCache: Tensor, + vSrc: Tensor, vCache: Tensor, + positions: Tensor, t: Int, + nKVHeads: Int, headDim: Int, maxSeq: Int, + on cmd: MTLCommandBuffer + ) { + precondition( + kSrc.dtype == kCache.dtype && vSrc.dtype == vCache.dtype + && kSrc.dtype == vSrc.dtype, + "Ops.kvCacheUpdateKVMany: dtype mismatch") + precondition( + positions.dtype == .u32, + "Ops.kvCacheUpdateKVMany: positions must be .u32") + precondition( + positions.elementCount == t, + "Ops.kvCacheUpdateKVMany: positions count must equal T") + let elementsPerRow = nKVHeads * headDim + let total = t * elementsPerRow + let (grid, tg) = elementwiseGrid(total) + let psoName: String + switch kSrc.dtype { + case .f32: psoName = "kv_cache_update_many_f32" + case .f16: psoName = "kv_cache_update_many_f16" + case .bf16: psoName = "kv_cache_update_many_bf16" + default: fatalError("Ops.kvCacheUpdateKVMany: unsupported dtype \(kSrc.dtype)") + } + let pso = PSOCache.shared.pipelineState(for: psoName) + guard let enc = cmd.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(pso) + // Shared bindings: positions + scalars set ONCE. + enc.setBuffer(positions.buffer, offset: positions.offset, index: 1) + var hd = UInt32(headDim) + var ms = UInt32(maxSeq) + var nhd = UInt32(elementsPerRow) + enc.setBytes(&hd, length: 4, index: 3) + enc.setBytes(&ms, length: 4, index: 4) + enc.setBytes(&nhd, length: 4, index: 5) + @inline(__always) + func dispatch(_ src: Tensor, _ cache: Tensor) { + enc.setBuffer(src.buffer, offset: src.offset, index: 0) + enc.setBuffer(cache.buffer, offset: cache.offset, index: 2) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + } + dispatch(kSrc, kCache) + dispatch(vSrc, vCache) + enc.endEncoding() + } + /// SDPA decode. q: [n_q_heads, head_dim]. k/v cache layout: /// [n_kv_heads, kv_stride, head_dim] where kv_stride is the physical /// capacity (maxSeq) and nKV is how many positions to attend to. @@ -2174,6 +3641,138 @@ public enum Ops { return result } + /// Two-pass Flash-Decoding SDPA. Pass 1 partitions the KV cache + /// into `blocks` even tiles and writes a per-block partial + /// `(O, max, lse)` triple; pass 2 merges them with the log-sum-exp + /// trick. Wins over single-pass `sdpaDecode` at long KV (≥ 1K + /// typical), where the single-pass kernel starves the GPU with + /// one threadgroup per Q-head. + /// + /// **Block-count contract.** The pass2 kernel hardcodes a 32-wide + /// reduction (`bn = 32` per lane) over `blocks / 32` chunks, so + /// `blocks` MUST be a multiple of 32 and at least 32. Passing + /// `blocks < 32` causes pass2 to iterate zero chunks and emit an + /// all-zero output silently — the precondition below catches that. + /// 32 is the standard short-context tile; 64 / 96 / 128 are + /// reasonable for very long contexts. + /// + /// **Scratch buffer sizing.** Caller owns the partials; reuse + /// across decode steps is safe because the shape depends only on + /// `(nQHeads, headDim, blocks)`: + /// - `partialO` : `[nQHeads, blocks, headDim]` matching `q.dtype` + /// - `partialM` : `[nQHeads, blocks]` fp32 (running max) + /// - `partialL` : `[nQHeads, blocks]` fp32 (running lse) + public static func sdpaDecode2Pass( + q: Tensor, k: Tensor, v: Tensor, + nQHeads: Int, nKVHeads: Int, headDim: Int, + nKV: Int, kvStride: Int, blocks: Int, + scale: Float, + partialO: Tensor, partialM: Tensor, partialL: Tensor, + into out: Tensor, + on cmd: MTLCommandBuffer + ) { + precondition( + nQHeads % nKVHeads == 0, + "Ops.sdpaDecode2Pass: nQHeads must be a multiple of nKVHeads") + let gqaFactor = nQHeads / nKVHeads + precondition( + blocks >= 32 && blocks % 32 == 0, + "Ops.sdpaDecode2Pass: blocks (\(blocks)) must be a multiple of 32 " + + "and ≥ 32 — pass2 hardcodes a 32-lane block-merge reduction") + precondition( + partialO.elementCount == nQHeads * blocks * headDim, + "Ops.sdpaDecode2Pass: partialO must be [nQHeads, blocks, headDim]") + precondition( + partialM.elementCount == nQHeads * blocks && partialM.dtype == .f32, + "Ops.sdpaDecode2Pass: partialM must be [nQHeads, blocks] f32") + precondition( + partialL.elementCount == nQHeads * blocks && partialL.dtype == .f32, + "Ops.sdpaDecode2Pass: partialL must be [nQHeads, blocks] f32") + precondition( + partialO.dtype == q.dtype && out.dtype == q.dtype, + "Ops.sdpaDecode2Pass: partialO / out dtype must match q") + + // Pass 1: per-block partial O / m / l. One TG per + // (kv_head, block) with `gqa_factor` simdgroups × 32 lanes, + // each lane owning `headDim / 32` consecutive elements. + let pass1TgWidth = gqaFactor * 32 + let pass1Grid = MTLSize( + width: nKVHeads * pass1TgWidth, height: blocks, depth: 1) + let pass1Tg = MTLSize(width: pass1TgWidth, height: 1, depth: 1) + + // Pass 2: merge partials across blocks. One TG per Q-head with + // 32 simdgroups × 32 lanes = 1024 threads. + let pass2TgWidth = 1024 + let pass2Grid = MTLSize( + width: nQHeads * pass2TgWidth, height: 1, depth: 1) + let pass2Tg = MTLSize(width: pass2TgWidth, height: 1, depth: 1) + + switch q.dtype { + case .f32: + MetalTileKernels.sdpa_decode_2pass_pass1_f32( + q: q.buffer, qOffset: q.offset, + k: k.buffer, kOffset: k.offset, + v: v.buffer, vOffset: v.offset, + partial_o: partialO.buffer, partial_oOffset: partialO.offset, + partial_m: partialM.buffer, partial_mOffset: partialM.offset, + partial_l: partialL.buffer, partial_lOffset: partialL.offset, + head_dim: UInt32(headDim), n_kv: UInt32(nKV), + kv_stride: UInt32(kvStride), + gqa_factor: UInt32(gqaFactor), blocks: UInt32(blocks), + scale: scale, + gridSize: pass1Grid, threadgroupSize: pass1Tg, on: cmd) + MetalTileKernels.sdpa_decode_2pass_pass2_f32( + partial_o: partialO.buffer, partial_oOffset: partialO.offset, + partial_m: partialM.buffer, partial_mOffset: partialM.offset, + partial_l: partialL.buffer, partial_lOffset: partialL.offset, + out: out.buffer, outOffset: out.offset, + head_dim: UInt32(headDim), blocks: UInt32(blocks), + gridSize: pass2Grid, threadgroupSize: pass2Tg, on: cmd) + case .f16: + MetalTileKernels.sdpa_decode_2pass_pass1_f16( + q: q.buffer, qOffset: q.offset, + k: k.buffer, kOffset: k.offset, + v: v.buffer, vOffset: v.offset, + partial_o: partialO.buffer, partial_oOffset: partialO.offset, + partial_m: partialM.buffer, partial_mOffset: partialM.offset, + partial_l: partialL.buffer, partial_lOffset: partialL.offset, + head_dim: UInt32(headDim), n_kv: UInt32(nKV), + kv_stride: UInt32(kvStride), + gqa_factor: UInt32(gqaFactor), blocks: UInt32(blocks), + scale: scale, + gridSize: pass1Grid, threadgroupSize: pass1Tg, on: cmd) + MetalTileKernels.sdpa_decode_2pass_pass2_f16( + partial_o: partialO.buffer, partial_oOffset: partialO.offset, + partial_m: partialM.buffer, partial_mOffset: partialM.offset, + partial_l: partialL.buffer, partial_lOffset: partialL.offset, + out: out.buffer, outOffset: out.offset, + head_dim: UInt32(headDim), blocks: UInt32(blocks), + gridSize: pass2Grid, threadgroupSize: pass2Tg, on: cmd) + case .bf16: + MetalTileKernels.sdpa_decode_2pass_pass1_bf16( + q: q.buffer, qOffset: q.offset, + k: k.buffer, kOffset: k.offset, + v: v.buffer, vOffset: v.offset, + partial_o: partialO.buffer, partial_oOffset: partialO.offset, + partial_m: partialM.buffer, partial_mOffset: partialM.offset, + partial_l: partialL.buffer, partial_lOffset: partialL.offset, + head_dim: UInt32(headDim), n_kv: UInt32(nKV), + kv_stride: UInt32(kvStride), + gqa_factor: UInt32(gqaFactor), blocks: UInt32(blocks), + scale: scale, + gridSize: pass1Grid, threadgroupSize: pass1Tg, on: cmd) + MetalTileKernels.sdpa_decode_2pass_pass2_bf16( + partial_o: partialO.buffer, partial_oOffset: partialO.offset, + partial_m: partialM.buffer, partial_mOffset: partialM.offset, + partial_l: partialL.buffer, partial_lOffset: partialL.offset, + out: out.buffer, outOffset: out.offset, + head_dim: UInt32(headDim), blocks: UInt32(blocks), + gridSize: pass2Grid, threadgroupSize: pass2Tg, on: cmd) + default: + fatalError("Ops.sdpaDecode2Pass: unsupported dtype \(q.dtype)") + } + } + /// Multi-query SDPA — attends `nQuery` query rows against a shared /// K/V cache in one dispatch. `q` / output are `[nQuery, nQHeads, /// headDim]`; `k` / `v` are the cache buffers `[nKVHeads, kvStride, @@ -2214,8 +3813,15 @@ public enum Ops { let grid = MTLSize(width: nQHeads * nQuery * threadsPerGroup, height: 1, depth: 1) let tg = MTLSize(width: threadsPerGroup, height: 1, depth: 1) let causalFlag = UInt32(causal ? 1 : 0) - switch q.dtype { - case .f32: + // Route by (dtype, headDim). d=128 → `ffai_sdpa_multi_*`; d=256 + // → `ffai_sdpa_multi_d256_*`, the variant added for Qwen3.6-A3B + // full-attention layers (head_dim=256). Both kernels share the + // same Swift wrapper shape; the dispatch grid / TPG / online + // softmax / causal mask / GQA fan-out semantics are identical. + // The d=256 kernel uses a 2-phase output reduction internally + // to stay under Apple's 32 KB threadgroup-memory cap. + switch (q.dtype, headDim) { + case (.f32, 128): MetalTileKernels.ffai_sdpa_multi_f32( q: q.buffer, qOffset: q.offset, k: k.buffer, kOffset: k.offset, v: v.buffer, vOffset: v.offset, out: result.buffer, outOffset: result.offset, @@ -2224,7 +3830,7 @@ public enum Ops { kv_stride: UInt32(kvStride), heads_per_group: UInt32(headsPerGroup), causal: causalFlag, scale: scale, gridSize: grid, threadgroupSize: tg, on: cmd) - case .f16: + case (.f16, 128): MetalTileKernels.ffai_sdpa_multi_f16( q: q.buffer, qOffset: q.offset, k: k.buffer, kOffset: k.offset, v: v.buffer, vOffset: v.offset, out: result.buffer, outOffset: result.offset, @@ -2233,7 +3839,7 @@ public enum Ops { kv_stride: UInt32(kvStride), heads_per_group: UInt32(headsPerGroup), causal: causalFlag, scale: scale, gridSize: grid, threadgroupSize: tg, on: cmd) - case .bf16: + case (.bf16, 128): MetalTileKernels.ffai_sdpa_multi_bf16( q: q.buffer, qOffset: q.offset, k: k.buffer, kOffset: k.offset, v: v.buffer, vOffset: v.offset, out: result.buffer, outOffset: result.offset, @@ -2242,8 +3848,123 @@ public enum Ops { kv_stride: UInt32(kvStride), heads_per_group: UInt32(headsPerGroup), causal: causalFlag, scale: scale, gridSize: grid, threadgroupSize: tg, on: cmd) + case (.f32, 256): + MetalTileKernels.ffai_sdpa_multi_d256_f32( + q: q.buffer, qOffset: q.offset, k: k.buffer, kOffset: k.offset, + v: v.buffer, vOffset: v.offset, out: result.buffer, outOffset: result.offset, + head_dim: UInt32(headDim), n_q_heads: UInt32(nQHeads), + base_kv: UInt32(baseKV), n_query: UInt32(nQuery), + kv_stride: UInt32(kvStride), heads_per_group: UInt32(headsPerGroup), + causal: causalFlag, scale: scale, + gridSize: grid, threadgroupSize: tg, on: cmd) + case (.f16, 256): + MetalTileKernels.ffai_sdpa_multi_d256_f16( + q: q.buffer, qOffset: q.offset, k: k.buffer, kOffset: k.offset, + v: v.buffer, vOffset: v.offset, out: result.buffer, outOffset: result.offset, + head_dim: UInt32(headDim), n_q_heads: UInt32(nQHeads), + base_kv: UInt32(baseKV), n_query: UInt32(nQuery), + kv_stride: UInt32(kvStride), heads_per_group: UInt32(headsPerGroup), + causal: causalFlag, scale: scale, + gridSize: grid, threadgroupSize: tg, on: cmd) + case (.bf16, 256): + MetalTileKernels.ffai_sdpa_multi_d256_bf16( + q: q.buffer, qOffset: q.offset, k: k.buffer, kOffset: k.offset, + v: v.buffer, vOffset: v.offset, out: result.buffer, outOffset: result.offset, + head_dim: UInt32(headDim), n_q_heads: UInt32(nQHeads), + base_kv: UInt32(baseKV), n_query: UInt32(nQuery), + kv_stride: UInt32(kvStride), heads_per_group: UInt32(headsPerGroup), + causal: causalFlag, scale: scale, + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError( + "Ops.sdpaMulti: unsupported dtype \(q.dtype) at headDim \(headDim) " + + "— only (f32/f16/bf16) × (128, 256) are emitted") + } + return result + } + + /// Tree-causal `sdpaMulti` — same dispatch contract but the + /// scalar `causal` boolean is replaced by an additive + /// `[nQuery, nQuery]` mask tensor consulted only for in-block KV + /// positions (positions `< baseKV` — the cached prefix — are + /// always fully attended). The mask is `0.0` for "allow" and + /// `-inf` for "block"; the kernel adds it directly to the + /// pre-softmax scores. + /// + /// Used by speculative-decode tree-verify: one verifier call + /// attends every leaf-to-root path of the draft tree at once, + /// gated by a tree-causal mask that lets each leaf only see its + /// own ancestors. Companion to `DraftTreeNode.treeCausalMask()`. + /// + /// Mask dtype must match `q`. `headDim` shares `sdpaMulti`'s + /// constraint: only the 128-variant kernel is emitted on dev + /// (tree-verify has no d256 use case today). + public static func sdpaMultiTreeMask( + q: Tensor, k: Tensor, v: Tensor, mask: Tensor, + nQHeads: Int, nKVHeads: Int, headDim: Int, + baseKV: Int, nQuery: Int, kvStride: Int, + scale: Float, + on cmd: MTLCommandBuffer, + into out: Tensor? = nil + ) -> Tensor { + if let reason = OpsValidation.validateSdpaMulti( + headDim: headDim, nQHeads: nQHeads, nKVHeads: nKVHeads, + baseKV: baseKV, nQuery: nQuery, kvStride: kvStride + ) { + preconditionFailure("Ops.sdpaMultiTreeMask: \(reason)") + } + precondition( + headDim == 128, + "Ops.sdpaMultiTreeMask: only headDim=128 has a tree-mask kernel variant") + precondition( + mask.dtype == q.dtype, + "Ops.sdpaMultiTreeMask: mask dtype (\(mask.dtype)) must match q dtype (\(q.dtype))") + precondition( + mask.elementCount == nQuery * nQuery, + "Ops.sdpaMultiTreeMask: mask elementCount \(mask.elementCount) ≠ nQuery·nQuery \(nQuery * nQuery)") + let headsPerGroup = nQHeads / nKVHeads + let result = out ?? Tensor.empty(shape: [nQuery, nQHeads, headDim], dtype: q.dtype) + let threadsPerGroup = 1024 + let grid = MTLSize( + width: nQHeads * nQuery * threadsPerGroup, + height: 1, depth: 1) + let tg = MTLSize(width: threadsPerGroup, height: 1, depth: 1) + switch q.dtype { + case .f32: + MetalTileKernels.ffai_sdpa_multi_tree_mask_f32( + q: q.buffer, qOffset: q.offset, k: k.buffer, kOffset: k.offset, + v: v.buffer, vOffset: v.offset, + mask: mask.buffer, maskOffset: mask.offset, + out: result.buffer, outOffset: result.offset, + head_dim: UInt32(headDim), n_q_heads: UInt32(nQHeads), + base_kv: UInt32(baseKV), n_query: UInt32(nQuery), + kv_stride: UInt32(kvStride), heads_per_group: UInt32(headsPerGroup), + scale: scale, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.ffai_sdpa_multi_tree_mask_f16( + q: q.buffer, qOffset: q.offset, k: k.buffer, kOffset: k.offset, + v: v.buffer, vOffset: v.offset, + mask: mask.buffer, maskOffset: mask.offset, + out: result.buffer, outOffset: result.offset, + head_dim: UInt32(headDim), n_q_heads: UInt32(nQHeads), + base_kv: UInt32(baseKV), n_query: UInt32(nQuery), + kv_stride: UInt32(kvStride), heads_per_group: UInt32(headsPerGroup), + scale: scale, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.ffai_sdpa_multi_tree_mask_bf16( + q: q.buffer, qOffset: q.offset, k: k.buffer, kOffset: k.offset, + v: v.buffer, vOffset: v.offset, + mask: mask.buffer, maskOffset: mask.offset, + out: result.buffer, outOffset: result.offset, + head_dim: UInt32(headDim), n_q_heads: UInt32(nQHeads), + base_kv: UInt32(baseKV), n_query: UInt32(nQuery), + kv_stride: UInt32(kvStride), heads_per_group: UInt32(headsPerGroup), + scale: scale, + gridSize: grid, threadgroupSize: tg, on: cmd) default: - fatalError("Ops.sdpaMulti: unsupported dtype \(q.dtype)") + fatalError("Ops.sdpaMultiTreeMask: unsupported dtype \(q.dtype)") } return result } @@ -3410,38 +5131,147 @@ public enum Ops { group_size: groupSizeU, gridSize: grid, threadgroupSize: tg, on: cmd) case .bf16: - MetalTileKernels.mt_moe_gather_qmm_mma_int4_bm8_mpp_bf16( - x: input.buffer, xOffset: input.offset, - w: weight.buffer, wOffset: weight.offset, - scales: scales.buffer, scalesOffset: scales.offset, - biases: biases.buffer, biasesOffset: biases.offset, - indices: indices.buffer, indicesOffset: indices.offset, - out: out.buffer, outOffset: out.offset, - m_total: mTotalU, n_out: nOutU, k_in: kInU, - group_size: groupSizeU, + MetalTileKernels.mt_moe_gather_qmm_mma_int4_bm8_mpp_bf16( + x: input.buffer, xOffset: input.offset, + w: weight.buffer, wOffset: weight.offset, + scales: scales.buffer, scalesOffset: scales.offset, + biases: biases.buffer, biasesOffset: biases.offset, + indices: indices.buffer, indicesOffset: indices.offset, + out: out.buffer, outOffset: out.offset, + m_total: mTotalU, n_out: nOutU, k_in: kInU, + group_size: groupSizeU, + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError("Ops.moeGatherDequantGemmInt4Bm8: unsupported dtype \(input.dtype)") + } + } + + // ─── Fused GDN prep + recurrence step ───────────────────────────── + // + // One dispatch absorbs the per-head q/k RMSNorm + g/beta math + + // the existing GDN recurrence step, collapsing the 3 host + // commit+wait pairs in Qwen35GDNMixer.forward down to 1. Same + // dispatch geometry as `mt_gated_delta_step`: + // grid = [dv, B·hv, 1] + // tg = [32, 1, 1] + // `convOut` layout: `[B, 2·Hk·Dk + Hv·Dv]` — q | k | v slabs. + // `qNormWeight` / `kNormWeight`: `[Hk·Dk]` (pass 1.0 × invKeyScale + // for the unweighted path). + public static func gatedDeltaPrepStep( + convOut: Tensor, + aLog: Tensor, dtBias: Tensor, + aRaw: Tensor, bRaw: Tensor, + qNormWeight: Tensor, kNormWeight: Tensor, + stateIn: Tensor, stateOut: Tensor, y: Tensor, + batchSize: Int, dk: Int, dv: Int, hv: Int, hk: Int, + on cmd: MTLCommandBuffer + ) { + precondition( + convOut.dtype == aLog.dtype && aLog.dtype == dtBias.dtype + && dtBias.dtype == aRaw.dtype && aRaw.dtype == bRaw.dtype + && bRaw.dtype == qNormWeight.dtype && qNormWeight.dtype == kNormWeight.dtype + && kNormWeight.dtype == stateIn.dtype && stateIn.dtype == stateOut.dtype + && stateOut.dtype == y.dtype, + "Ops.gatedDeltaPrepStep: every tensor must share dtype") + precondition( + dk % 32 == 0, + "Ops.gatedDeltaPrepStep: dk \(dk) must be multiple of 32") + precondition( + dv % 32 == 0, + "Ops.gatedDeltaPrepStep: dv \(dv) must be multiple of 32") + precondition( + hv % hk == 0, + "Ops.gatedDeltaPrepStep: hv \(hv) must be a multiple of hk \(hk) (GQA)") + + // `dispatchThreads` counts TOTAL threads per axis, so the + // X axis is `dv` (TGs along X) × `tgWidth` (threads per TG). + // `tgid_x` inside the kernel therefore ranges 0..dv-1, matching + // the kernel's `dv_idx = tgid_x` contract. Y axis is one TG + // per (batch · Hv) slab. + let tgWidth = 32 + let grid = MTLSize( + width: dv * tgWidth, + height: batchSize * hv, + depth: 1) + let tg = MTLSize(width: tgWidth, height: 1, depth: 1) + let dkU = UInt32(dk) + let dvU = UInt32(dv) + let hvU = UInt32(hv) + let hkU = UInt32(hk) + + switch convOut.dtype { + case .f32: + MetalTileKernels.mt_gated_delta_prep_step_f32( + conv_out: convOut.buffer, conv_outOffset: convOut.offset, + a_log: aLog.buffer, a_logOffset: aLog.offset, + dt_bias: dtBias.buffer, dt_biasOffset: dtBias.offset, + a_raw: aRaw.buffer, a_rawOffset: aRaw.offset, + b_raw: bRaw.buffer, b_rawOffset: bRaw.offset, + q_norm_weight: qNormWeight.buffer, q_norm_weightOffset: qNormWeight.offset, + k_norm_weight: kNormWeight.buffer, k_norm_weightOffset: kNormWeight.offset, + state_in: stateIn.buffer, state_inOffset: stateIn.offset, + state_out: stateOut.buffer, state_outOffset: stateOut.offset, + y: y.buffer, yOffset: y.offset, + dk: dkU, dv: dvU, hv: hvU, hk: hkU, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.mt_gated_delta_prep_step_f16( + conv_out: convOut.buffer, conv_outOffset: convOut.offset, + a_log: aLog.buffer, a_logOffset: aLog.offset, + dt_bias: dtBias.buffer, dt_biasOffset: dtBias.offset, + a_raw: aRaw.buffer, a_rawOffset: aRaw.offset, + b_raw: bRaw.buffer, b_rawOffset: bRaw.offset, + q_norm_weight: qNormWeight.buffer, q_norm_weightOffset: qNormWeight.offset, + k_norm_weight: kNormWeight.buffer, k_norm_weightOffset: kNormWeight.offset, + state_in: stateIn.buffer, state_inOffset: stateIn.offset, + state_out: stateOut.buffer, state_outOffset: stateOut.offset, + y: y.buffer, yOffset: y.offset, + dk: dkU, dv: dvU, hv: hvU, hk: hkU, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.mt_gated_delta_prep_step_bf16( + conv_out: convOut.buffer, conv_outOffset: convOut.offset, + a_log: aLog.buffer, a_logOffset: aLog.offset, + dt_bias: dtBias.buffer, dt_biasOffset: dtBias.offset, + a_raw: aRaw.buffer, a_rawOffset: aRaw.offset, + b_raw: bRaw.buffer, b_rawOffset: bRaw.offset, + q_norm_weight: qNormWeight.buffer, q_norm_weightOffset: qNormWeight.offset, + k_norm_weight: kNormWeight.buffer, k_norm_weightOffset: kNormWeight.offset, + state_in: stateIn.buffer, state_inOffset: stateIn.offset, + state_out: stateOut.buffer, state_outOffset: stateOut.offset, + y: y.buffer, yOffset: y.offset, + dk: dkU, dv: dvU, hv: hvU, hk: hkU, gridSize: grid, threadgroupSize: tg, on: cmd) default: - fatalError("Ops.moeGatherDequantGemmInt4Bm8: unsupported dtype \(input.dtype)") + fatalError("Ops.gatedDeltaPrepStep: unsupported dtype \(convOut.dtype)") } } - // ─── Fused GDN prep + recurrence step ───────────────────────────── - // - // One dispatch absorbs the per-head q/k RMSNorm + g/beta math + - // the existing GDN recurrence step, collapsing the 3 host - // commit+wait pairs in Qwen35GDNMixer.forward down to 1. Same - // dispatch geometry as `mt_gated_delta_step`: - // grid = [dv, B·hv, 1] - // tg = [32, 1, 1] - // `convOut` layout: `[B, 2·Hk·Dk + Hv·Dv]` — q | k | v slabs. - // `qNormWeight` / `kNormWeight`: `[Hk·Dk]` (pass 1.0 × invKeyScale - // for the unweighted path). - public static func gatedDeltaPrepStep( + /// Chunked-prefill counterpart to `gatedDeltaPrepStep`. Runs the + /// fused GDN prep + recurrence over a contiguous range of tokens + /// in one dispatch instead of looping the single-token kernel, + /// keeping the per-head recurrent state in registers across the + /// chunk. `tLen` is a GPU-resident `[1]` u32 scalar so the host + /// can chain chunked dispatches without a CPU readback for the + /// chunk-size constant. + /// + /// Shape contract (all share dtype except `tLen`): + /// - `convOut` : `[batchSize, tLen, hv·dv]` + /// - `aLog` : `[hv]` per-head log-decay + /// - `dtBias` : `[hv]` per-head delta-t bias + /// - `aRaw`/`bRaw` : `[batchSize, tLen, hv]` chunk gates + /// - `qNormWeight` / `kNormWeight` : `[dk]` Q-norm / K-norm + /// applied before the recurrence + /// - `stateIn` / `stateOut` : `[batchSize, hv, dv, dk]` fp32 + /// recurrent state (read-once / write-once across the chunk) + /// - `y` : `[batchSize, tLen, hv·dv]` output activations + public static func gatedDeltaPrepChunk( convOut: Tensor, aLog: Tensor, dtBias: Tensor, aRaw: Tensor, bRaw: Tensor, qNormWeight: Tensor, kNormWeight: Tensor, stateIn: Tensor, stateOut: Tensor, y: Tensor, + tLen: Tensor, batchSize: Int, dk: Int, dv: Int, hv: Int, hk: Int, on cmd: MTLCommandBuffer ) { @@ -3451,22 +5281,21 @@ public enum Ops { && bRaw.dtype == qNormWeight.dtype && qNormWeight.dtype == kNormWeight.dtype && kNormWeight.dtype == stateIn.dtype && stateIn.dtype == stateOut.dtype && stateOut.dtype == y.dtype, - "Ops.gatedDeltaPrepStep: every tensor must share dtype") + "Ops.gatedDeltaPrepChunk: every tensor must share dtype") precondition( dk % 32 == 0, - "Ops.gatedDeltaPrepStep: dk \(dk) must be multiple of 32") + "Ops.gatedDeltaPrepChunk: dk (\(dk)) must be a multiple of 32") precondition( dv % 32 == 0, - "Ops.gatedDeltaPrepStep: dv \(dv) must be multiple of 32") + "Ops.gatedDeltaPrepChunk: dv (\(dv)) must be a multiple of 32") precondition( hv % hk == 0, - "Ops.gatedDeltaPrepStep: hv \(hv) must be a multiple of hk \(hk) (GQA)") - - // `dispatchThreads` counts TOTAL threads per axis, so the - // X axis is `dv` (TGs along X) × `tgWidth` (threads per TG). - // `tgid_x` inside the kernel therefore ranges 0..dv-1, matching - // the kernel's `dv_idx = tgid_x` contract. Y axis is one TG - // per (batch · Hv) slab. + "Ops.gatedDeltaPrepChunk: hv (\(hv)) must be a multiple of hk (\(hk)) (GQA)") + precondition( + tLen.dtype == .u32 && tLen.elementCount == 1, + "Ops.gatedDeltaPrepChunk: tLen must be a [1] u32 scalar buffer") + // One simdgroup per (head, dv index) ; grid sweeps batch · hv + // along Y. Kernel walks tLen tokens internally. let tgWidth = 32 let grid = MTLSize( width: dv * tgWidth, @@ -3477,10 +5306,9 @@ public enum Ops { let dvU = UInt32(dv) let hvU = UInt32(hv) let hkU = UInt32(hk) - switch convOut.dtype { case .f32: - MetalTileKernels.mt_gated_delta_prep_step_f32( + MetalTileKernels.mt_gated_delta_prep_chunk_f32( conv_out: convOut.buffer, conv_outOffset: convOut.offset, a_log: aLog.buffer, a_logOffset: aLog.offset, dt_bias: dtBias.buffer, dt_biasOffset: dtBias.offset, @@ -3491,10 +5319,11 @@ public enum Ops { state_in: stateIn.buffer, state_inOffset: stateIn.offset, state_out: stateOut.buffer, state_outOffset: stateOut.offset, y: y.buffer, yOffset: y.offset, + t_len: tLen.buffer, t_lenOffset: tLen.offset, dk: dkU, dv: dvU, hv: hvU, hk: hkU, gridSize: grid, threadgroupSize: tg, on: cmd) case .f16: - MetalTileKernels.mt_gated_delta_prep_step_f16( + MetalTileKernels.mt_gated_delta_prep_chunk_f16( conv_out: convOut.buffer, conv_outOffset: convOut.offset, a_log: aLog.buffer, a_logOffset: aLog.offset, dt_bias: dtBias.buffer, dt_biasOffset: dtBias.offset, @@ -3505,10 +5334,11 @@ public enum Ops { state_in: stateIn.buffer, state_inOffset: stateIn.offset, state_out: stateOut.buffer, state_outOffset: stateOut.offset, y: y.buffer, yOffset: y.offset, + t_len: tLen.buffer, t_lenOffset: tLen.offset, dk: dkU, dv: dvU, hv: hvU, hk: hkU, gridSize: grid, threadgroupSize: tg, on: cmd) case .bf16: - MetalTileKernels.mt_gated_delta_prep_step_bf16( + MetalTileKernels.mt_gated_delta_prep_chunk_bf16( conv_out: convOut.buffer, conv_outOffset: convOut.offset, a_log: aLog.buffer, a_logOffset: aLog.offset, dt_bias: dtBias.buffer, dt_biasOffset: dtBias.offset, @@ -3519,10 +5349,11 @@ public enum Ops { state_in: stateIn.buffer, state_inOffset: stateIn.offset, state_out: stateOut.buffer, state_outOffset: stateOut.offset, y: y.buffer, yOffset: y.offset, + t_len: tLen.buffer, t_lenOffset: tLen.offset, dk: dkU, dv: dvU, hv: hvU, hk: hkU, gridSize: grid, threadgroupSize: tg, on: cmd) default: - fatalError("Ops.gatedDeltaPrepStep: unsupported dtype \(convOut.dtype)") + fatalError("Ops.gatedDeltaPrepChunk: unsupported dtype \(convOut.dtype)") } } @@ -3743,6 +5574,103 @@ public enum Ops { } } + /// Cast TWO same-dtype tensors to f32 on the same compute encoder. + /// Saves one encoder begin/end pair versus two `castToF32` calls. + /// Used inside the Qwen3 GDN T-loop where the two raw gate inputs + /// (`aRaw`, `bRaw`) need to be promoted to fp32 every token. + public static func castToF32Two( + _ a: Tensor, into outA: Tensor, + _ b: Tensor, into outB: Tensor, + on cmd: MTLCommandBuffer + ) { + precondition( + a.dtype == b.dtype, + "Ops.castToF32Two: inputs must share dtype") + precondition( + outA.dtype == .f32 && outB.dtype == .f32, + "Ops.castToF32Two: outputs must be f32") + precondition( + a.elementCount == outA.elementCount, + "Ops.castToF32Two: a / outA element-count mismatch") + precondition( + b.elementCount == outB.elementCount, + "Ops.castToF32Two: b / outB element-count mismatch") + let psoName: String + switch a.dtype { + case .bf16: psoName = "mt_cast_to_f32_bf16" + case .f16: psoName = "mt_cast_to_f32_f16" + case .f32: psoName = "mt_cast_to_f32_f32" + default: fatalError("Ops.castToF32Two: unsupported dtype \(a.dtype)") + } + let pso = PSOCache.shared.pipelineState(for: psoName) + guard let enc = cmd.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(pso) + @inline(__always) + func dispatch(_ input: Tensor, _ out: Tensor) { + enc.setBuffer(input.buffer, offset: input.offset, index: 0) + enc.setBuffer(out.buffer, offset: out.offset, index: 1) + let n = input.elementCount + let tgWidth = min(n, 256) + enc.dispatchThreads( + MTLSize(width: n, height: 1, depth: 1), + threadsPerThreadgroup: MTLSize(width: tgWidth, height: 1, depth: 1)) + } + dispatch(a, outA) + dispatch(b, outB) + enc.endEncoding() + } + + /// Cast THREE same-dtype tensors to f32 on the same compute + /// encoder. Used in the Qwen3 GDN fused-prep path where + /// `convAct`, `aRaw`, and `bRaw` all need f32 promotion before + /// the recurrence kernel. + public static func castToF32Three( + _ a: Tensor, into outA: Tensor, + _ b: Tensor, into outB: Tensor, + _ c: Tensor, into outC: Tensor, + on cmd: MTLCommandBuffer + ) { + precondition( + a.dtype == b.dtype && b.dtype == c.dtype, + "Ops.castToF32Three: all inputs must share dtype") + precondition( + outA.dtype == .f32 && outB.dtype == .f32 && outC.dtype == .f32, + "Ops.castToF32Three: outputs must all be f32") + precondition( + a.elementCount == outA.elementCount, + "Ops.castToF32Three: a / outA element-count mismatch") + precondition( + b.elementCount == outB.elementCount, + "Ops.castToF32Three: b / outB element-count mismatch") + precondition( + c.elementCount == outC.elementCount, + "Ops.castToF32Three: c / outC element-count mismatch") + let psoName: String + switch a.dtype { + case .bf16: psoName = "mt_cast_to_f32_bf16" + case .f16: psoName = "mt_cast_to_f32_f16" + case .f32: psoName = "mt_cast_to_f32_f32" + default: fatalError("Ops.castToF32Three: unsupported dtype \(a.dtype)") + } + let pso = PSOCache.shared.pipelineState(for: psoName) + guard let enc = cmd.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(pso) + @inline(__always) + func dispatch(_ input: Tensor, _ out: Tensor) { + enc.setBuffer(input.buffer, offset: input.offset, index: 0) + enc.setBuffer(out.buffer, offset: out.offset, index: 1) + let n = input.elementCount + let tgWidth = min(n, 256) + enc.dispatchThreads( + MTLSize(width: n, height: 1, depth: 1), + threadsPerThreadgroup: MTLSize(width: tgWidth, height: 1, depth: 1)) + } + dispatch(a, outA) + dispatch(b, outB) + dispatch(c, outC) + enc.endEncoding() + } + public static func castToF32( _ input: Tensor, into output: Tensor, on cmd: MTLCommandBuffer @@ -3781,6 +5709,67 @@ public enum Ops { } } + /// One silu-cast plus two plain casts to f32 on the SAME compute + /// encoder. Used inside the Qwen3.5 GDN T-loop where, every token, + /// the input goes through `siluCastToF32` while `aRaw` / `bRaw` + /// (the gate scalars) go through plain `castToF32`. The wrapper + /// sets the silu PSO once, dispatches the silu, then switches to + /// the plain-cast PSO and dispatches the two raw casts. Saves the + /// encoder begin/end pair between dispatches that would otherwise + /// fire on every GDN token. + public static func siluCastF32PlusCastF32Two( + siluIn: Tensor, into siluOut: Tensor, + _ a: Tensor, into outA: Tensor, + _ b: Tensor, into outB: Tensor, + on cmd: MTLCommandBuffer + ) { + precondition( + siluIn.dtype == a.dtype && a.dtype == b.dtype, + "Ops.siluCastF32PlusCastF32Two: all inputs must share dtype") + precondition( + siluOut.dtype == .f32 && outA.dtype == .f32 && outB.dtype == .f32, + "Ops.siluCastF32PlusCastF32Two: outputs must be f32") + precondition( + siluIn.elementCount == siluOut.elementCount, + "Ops.siluCastF32PlusCastF32Two: silu in/out count mismatch") + precondition( + a.elementCount == outA.elementCount, + "Ops.siluCastF32PlusCastF32Two: a/outA count mismatch") + precondition( + b.elementCount == outB.elementCount, + "Ops.siluCastF32PlusCastF32Two: b/outB count mismatch") + let siluPso: String + let castPso: String + switch siluIn.dtype { + case .f16: + siluPso = "mt_silu_cast_to_f32_f16" + castPso = "mt_cast_to_f32_f16" + case .bf16: + siluPso = "mt_silu_cast_to_f32_bf16" + castPso = "mt_cast_to_f32_bf16" + default: + fatalError( + "Ops.siluCastF32PlusCastF32Two: unsupported dtype \(siluIn.dtype)") + } + guard let enc = cmd.makeComputeCommandEncoder() else { return } + @inline(__always) + func dispatch(_ input: Tensor, _ out: Tensor) { + enc.setBuffer(input.buffer, offset: input.offset, index: 0) + enc.setBuffer(out.buffer, offset: out.offset, index: 1) + let n = input.elementCount + let tgWidth = min(n, 256) + enc.dispatchThreads( + MTLSize(width: n, height: 1, depth: 1), + threadsPerThreadgroup: MTLSize(width: tgWidth, height: 1, depth: 1)) + } + enc.setComputePipelineState(PSOCache.shared.pipelineState(for: siluPso)) + dispatch(siluIn, siluOut) + enc.setComputePipelineState(PSOCache.shared.pipelineState(for: castPso)) + dispatch(a, outA) + dispatch(b, outB) + enc.endEncoding() + } + // ─── Fused gated mixer norm ─────────────────────────────────────── // // Wraps `mt_gated_mixer_norm_{f32,f16,bf16}`. Computes @@ -3874,6 +5863,309 @@ public enum Ops { } } + /// T-batched `gatedMixerNorm`. Same kernel, same per-row geometry + /// (one threadgroup per row, `Dv / 4` threads per TG); the grid X + /// axis widens from `Hv` to `T · Hv` so a single dispatch covers + /// every prefill row. Used in the Qwen3.5 GDN mixer's prefill + /// path where the T-loop over per-token norms was an explicit + /// performance hotspot. + /// + /// Tensor shapes: + /// - `y` `[T, Hv, Dv]` fp32 (recurrence output stays fp32) + /// - `z` / `out` `[T, Hv, Dv]` in model dtype + /// - `weight` `[Dv]` in model dtype + /// - `epsBuf` `[1]` fp32 + public static func gatedMixerNormMany( + y: Tensor, z: Tensor, weight: Tensor, + epsBuf: Tensor, + into out: Tensor, + t: Int, numValueHeads: Int, valueHeadDim: Int, + on cmd: MTLCommandBuffer + ) { + precondition(t > 0, "Ops.gatedMixerNormMany: t must be positive") + precondition( + y.dtype == .f32, + "Ops.gatedMixerNormMany: y must be f32 (got \(y.dtype))") + precondition( + z.dtype == weight.dtype && weight.dtype == out.dtype, + "Ops.gatedMixerNormMany: z / weight / out must share dtype") + precondition( + epsBuf.dtype == .f32, + "Ops.gatedMixerNormMany: epsBuf must be f32") + precondition( + valueHeadDim.isMultiple(of: 4), + "Ops.gatedMixerNormMany: valueHeadDim (\(valueHeadDim)) must be a multiple of 4") + let expected = t * numValueHeads * valueHeadDim + precondition( + y.elementCount == expected, + "Ops.gatedMixerNormMany: y has \(y.elementCount), expected T·Hv·Dv = \(expected)") + precondition( + z.elementCount == expected, + "Ops.gatedMixerNormMany: z has \(z.elementCount), expected \(expected)") + precondition( + weight.elementCount == valueHeadDim, + "Ops.gatedMixerNormMany: weight has \(weight.elementCount), expected Dv = \(valueHeadDim)") + precondition( + out.elementCount == expected, + "Ops.gatedMixerNormMany: out has \(out.elementCount), expected \(expected)") + let tpg = valueHeadDim / 4 + let nRows = t * numValueHeads + let grid = MTLSize(width: nRows * tpg, height: 1, depth: 1) + let tg = MTLSize(width: tpg, height: 1, depth: 1) + let nU = UInt32(valueHeadDim) + switch out.dtype { + case .f32: + MetalTileKernels.mt_gated_mixer_norm_f32( + y: y.buffer, yOffset: y.offset, + z: z.buffer, zOffset: z.offset, + w: weight.buffer, wOffset: weight.offset, + out: out.buffer, outOffset: out.offset, + eps_buf: epsBuf.buffer, eps_bufOffset: epsBuf.offset, + n: nU, gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.mt_gated_mixer_norm_f16( + y: y.buffer, yOffset: y.offset, + z: z.buffer, zOffset: z.offset, + w: weight.buffer, wOffset: weight.offset, + out: out.buffer, outOffset: out.offset, + eps_buf: epsBuf.buffer, eps_bufOffset: epsBuf.offset, + n: nU, gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.mt_gated_mixer_norm_bf16( + y: y.buffer, yOffset: y.offset, + z: z.buffer, zOffset: z.offset, + w: weight.buffer, wOffset: weight.offset, + out: out.buffer, outOffset: out.offset, + eps_buf: epsBuf.buffer, eps_bufOffset: epsBuf.offset, + n: nU, gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError("Ops.gatedMixerNormMany: unsupported out dtype \(out.dtype)") + } + } + + // ─── Fused RMSNorm + int4 dequant-GEMV ──────────────────────────── + // + // Wraps `ffai_rms_norm_qgemv_fast` (RMSNorm-fused) and + // `ffai_gated_rms_norm_qgemv_int4_fast` (Mamba2-gated-RMSNorm- + // fused). Both keep the normalised activation in registers and + // feed it straight into the int4 dequant-GEMV, replacing a + // 2-dispatch chain (`Ops.rmsNorm` / `Ops.gatedRmsNorm` + `Ops. + // dequantGemvInt4`) with a single kernel. Saves a full `[in_dim]` + // DRAM roundtrip on the intermediate `normed` per call. + // + // Use at every site where a pre-norm immediately precedes a + // single int4 qmm projection — most notably the finalNorm + lmHead + // boundary (one dispatch per token) and the GDN mixer's + // post-recurrence gatedRmsNorm + outProj pair. + + /// Fused RMSNorm + int4 dequant-GEMV in ONE dispatch via + /// `ffai_rms_norm_qgemv_fast`. Computes + /// `y[row] = Σ_i (q[row,i]·scale + bias) · + /// (x[i] · norm_weight[i] · inv_rms)` + /// with `inv_rms = rsqrt(mean(x²) + eps)` and the normalised + /// activation never leaving registers. + /// + /// Kernel constraints (`ffai_rms_norm_qgemv_fast`): + /// - `in_dim` MUST be a multiple of 512 (kernel block size = 512 + /// K-elements per outer iter). + /// - `out_dim` MUST be a multiple of 8 (kernel processes 8 output + /// rows per TG). + /// - `group_size` MUST equal 64. + /// - TPG = 64 (2 simdgroups × 32 lanes). + public static func rmsNormQgemvInt4Fast( + x: Tensor, normWeight: Tensor, eps: Float, + qWeight: Tensor, qScales: Tensor, qBiases: Tensor, + on cmd: MTLCommandBuffer, + into out: Tensor + ) { + precondition( + qWeight.dtype == .u32, + "Ops.rmsNormQgemvInt4Fast: qWeight must be u32-packed") + precondition( + qWeight.shape.count == 2, + "Ops.rmsNormQgemvInt4Fast: qWeight must be [outDim, inDim/8]") + let outDim = qWeight.shape[0] + let packedPerRow = qWeight.shape[1] + let inDim = packedPerRow * 8 + precondition( + x.elementCount == inDim, + "Ops.rmsNormQgemvInt4Fast: x.elementCount \(x.elementCount) ≠ inDim \(inDim)") + precondition( + normWeight.elementCount == inDim, + "Ops.rmsNormQgemvInt4Fast: normWeight.elementCount \(normWeight.elementCount) ≠ inDim \(inDim)") + precondition( + out.elementCount == outDim, + "Ops.rmsNormQgemvInt4Fast: out.elementCount \(out.elementCount) ≠ outDim \(outDim)") + precondition( + x.dtype == normWeight.dtype && normWeight.dtype == qScales.dtype + && qScales.dtype == qBiases.dtype && qBiases.dtype == out.dtype, + "Ops.rmsNormQgemvInt4Fast: all non-weight tensors must share dtype") + precondition( + inDim % 512 == 0, + "Ops.rmsNormQgemvInt4Fast: in_dim \(inDim) must be a multiple of 512 (fast variant)") + precondition( + outDim % 8 == 0, + "Ops.rmsNormQgemvInt4Fast: out_dim \(outDim) must be a multiple of 8") + // INVARIANT: kernel pins group_size=64 + TPG=64 for the fast + // 8-rows-per-TG variant. group_size is checked as a constexpr + // by the kernel; we re-assert here for early failure. + let groupSize = 64 + precondition( + inDim % groupSize == 0, + "Ops.rmsNormQgemvInt4Fast: in_dim must divide group_size=64") + // eps as a 1-element f32 buffer. + var epsValue = eps + let epsBuf = device.makeBuffer(length: 4) + memcpy(epsBuf.contents(), &epsValue, 4) + let tg = MTLSize(width: 64, height: 1, depth: 1) + let nTiles = outDim / 8 + let grid = MTLSize(width: nTiles * 64, height: 1, depth: 1) + switch x.dtype { + case .f32: + MetalTileKernels.ffai_rms_norm_qgemv_fast_f32( + x: x.buffer, xOffset: x.offset, + norm_weight: normWeight.buffer, norm_weightOffset: normWeight.offset, + weight: qWeight.buffer, weightOffset: qWeight.offset, + scales: qScales.buffer, scalesOffset: qScales.offset, + biases: qBiases.buffer, biasesOffset: qBiases.offset, + output: out.buffer, outputOffset: out.offset, + eps_buf: epsBuf, eps_bufOffset: 0, + in_dim: UInt32(inDim), group_size: UInt32(groupSize), + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.ffai_rms_norm_qgemv_fast_f16( + x: x.buffer, xOffset: x.offset, + norm_weight: normWeight.buffer, norm_weightOffset: normWeight.offset, + weight: qWeight.buffer, weightOffset: qWeight.offset, + scales: qScales.buffer, scalesOffset: qScales.offset, + biases: qBiases.buffer, biasesOffset: qBiases.offset, + output: out.buffer, outputOffset: out.offset, + eps_buf: epsBuf, eps_bufOffset: 0, + in_dim: UInt32(inDim), group_size: UInt32(groupSize), + gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.ffai_rms_norm_qgemv_fast_bf16( + x: x.buffer, xOffset: x.offset, + norm_weight: normWeight.buffer, norm_weightOffset: normWeight.offset, + weight: qWeight.buffer, weightOffset: qWeight.offset, + scales: qScales.buffer, scalesOffset: qScales.offset, + biases: qBiases.buffer, biasesOffset: qBiases.offset, + output: out.buffer, outputOffset: out.offset, + eps_buf: epsBuf, eps_bufOffset: 0, + in_dim: UInt32(inDim), group_size: UInt32(groupSize), + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError("Ops.rmsNormQgemvInt4Fast: unsupported dtype \(x.dtype)") + } + } + + /// Fused Mamba2-style gated RMSNorm + int4 dequant-GEMV in ONE + /// dispatch via `ffai_gated_rms_norm_qgemv_int4_fast`. Computes + /// the GDN mixer's `silu(z) · w · rmsnorm(y)` and feeds the result + /// straight into the int4 outProj — replaces the 2-dispatch + /// chain (`Ops.gatedRmsNorm` + `Ops.dequantGemvInt4`) with a + /// single kernel. Saves a full `[Hv·Dv]` DRAM roundtrip on the + /// intermediate gated-norm output per call. + /// + /// Use at the GDN mixer post-recurrence boundary, where the + /// gated-norm output is the direct input to the outProj projection. + /// + /// Kernel constraints: + /// - `Hv·Dv` (in_dim) MUST be a multiple of 512. + /// - `Hv·Dv` MUST be ≤ 8192 (kernel TG-memory cap; the gated-norm + /// pass stages the full input in TG memory). + /// - `out_dim` MUST be a multiple of 8. + /// - `group_size` MUST be 64. TPG = 64. + public static func gatedRmsNormQgemvInt4Fast( + y: Tensor, // [Hv, Dv] f32 (GDN recurrence output) + z: Tensor, // [Hv*Dv] T + normWeight: Tensor, // [Dv] T + eps: Float, + qWeight: Tensor, // [out_dim, in_dim/8] u32 + qScales: Tensor, // [out_dim, in_dim/group_size] T + qBiases: Tensor, + hv: Int, dv: Int, outDim: Int, groupSize: Int = 64, + on cmd: MTLCommandBuffer, + into out: Tensor // [out_dim] T + ) { + precondition( + y.dtype == .f32, + "Ops.gatedRmsNormQgemvInt4Fast: y must be f32") + precondition( + z.dtype == normWeight.dtype && normWeight.dtype == out.dtype, + "Ops.gatedRmsNormQgemvInt4Fast: z/weight/out dtype mismatch") + precondition( + qWeight.dtype == .u32, + "Ops.gatedRmsNormQgemvInt4Fast: q_weight must be u32-packed") + let inDim = hv * dv + precondition( + inDim % 512 == 0, + "Ops.gatedRmsNormQgemvInt4Fast: Hv·Dv (\(inDim)) must be multiple of 512") + precondition( + inDim <= 8192, + "Ops.gatedRmsNormQgemvInt4Fast: Hv·Dv (\(inDim)) must be ≤ 8192 (kernel TG-mem cap)") + precondition( + outDim % 8 == 0, + "Ops.gatedRmsNormQgemvInt4Fast: out_dim must be multiple of 8") + precondition( + groupSize == 64, + "Ops.gatedRmsNormQgemvInt4Fast: group_size must be 64") + precondition( + out.elementCount == outDim, + "Ops.gatedRmsNormQgemvInt4Fast: out element count mismatch") + let tpg = 64 + let nTiles = outDim / 8 + let grid = MTLSize(width: nTiles * tpg, height: 1, depth: 1) + let tg = MTLSize(width: tpg, height: 1, depth: 1) + // eps as a 1-element f32 buffer. + var epsValue = eps + let epsBuf = device.makeBuffer(length: 4) + memcpy(epsBuf.contents(), &epsValue, 4) + switch out.dtype { + case .f32: + MetalTileKernels.ffai_gated_rms_norm_qgemv_int4_fast_f32( + y: y.buffer, yOffset: y.offset, + z: z.buffer, zOffset: z.offset, + norm_weight: normWeight.buffer, norm_weightOffset: normWeight.offset, + eps_buf: epsBuf, eps_bufOffset: 0, + q_weight: qWeight.buffer, q_weightOffset: qWeight.offset, + q_scales: qScales.buffer, q_scalesOffset: qScales.offset, + q_biases: qBiases.buffer, q_biasesOffset: qBiases.offset, + out: out.buffer, outOffset: out.offset, + hv: UInt32(hv), dv: UInt32(dv), + out_dim: UInt32(outDim), group_size: UInt32(groupSize), + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.ffai_gated_rms_norm_qgemv_int4_fast_f16( + y: y.buffer, yOffset: y.offset, + z: z.buffer, zOffset: z.offset, + norm_weight: normWeight.buffer, norm_weightOffset: normWeight.offset, + eps_buf: epsBuf, eps_bufOffset: 0, + q_weight: qWeight.buffer, q_weightOffset: qWeight.offset, + q_scales: qScales.buffer, q_scalesOffset: qScales.offset, + q_biases: qBiases.buffer, q_biasesOffset: qBiases.offset, + out: out.buffer, outOffset: out.offset, + hv: UInt32(hv), dv: UInt32(dv), + out_dim: UInt32(outDim), group_size: UInt32(groupSize), + gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.ffai_gated_rms_norm_qgemv_int4_fast_bf16( + y: y.buffer, yOffset: y.offset, + z: z.buffer, zOffset: z.offset, + norm_weight: normWeight.buffer, norm_weightOffset: normWeight.offset, + eps_buf: epsBuf, eps_bufOffset: 0, + q_weight: qWeight.buffer, q_weightOffset: qWeight.offset, + q_scales: qScales.buffer, q_scalesOffset: qScales.offset, + q_biases: qBiases.buffer, q_biasesOffset: qBiases.offset, + out: out.buffer, outOffset: out.offset, + hv: UInt32(hv), dv: UInt32(dv), + out_dim: UInt32(outDim), group_size: UInt32(groupSize), + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError("Ops.gatedRmsNormQgemvInt4Fast: unsupported dtype \(out.dtype)") + } + } + // ─── Fused scalar-sigmoid fan-out + FMA ─────────────────────────── // // Wraps `mt_sigmoid_scalar_fma_{f32,f16,bf16}`. Computes @@ -4072,6 +6364,51 @@ public enum Ops { return result } + /// N independent SwiGLU dispatches sharing ONE compute encoder. + /// Each `(gate, up, out)` triple is dispatched in sequence after + /// setting the same `mt_swiglu_*` PSO once. Used by the MoE + /// per-expert SwiGLU phase at decode (one dispatch per chosen + /// expert × topK experts × MoE layers) where the encoder + /// begin/end pairs dominated CPU side. + public static func swigluMany( + gates: [Tensor], ups: [Tensor], outs: [Tensor], + on cmd: MTLCommandBuffer + ) { + let n = gates.count + precondition( + ups.count == n && outs.count == n, + "Ops.swigluMany: count mismatch") + guard n > 0 else { return } + let dtype = gates[0].dtype + let psoName: String + switch dtype { + case .f32: psoName = "mt_swiglu_f32" + case .f16: psoName = "mt_swiglu_f16" + case .bf16: psoName = "mt_swiglu_bf16" + default: fatalError("Ops.swigluMany: unsupported dtype \(dtype)") + } + let pso = PSOCache.shared.pipelineState(for: psoName) + guard let enc = cmd.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(pso) + for i in 0 ..< n { + let count = gates[i].elementCount + precondition( + ups[i].elementCount == count && outs[i].elementCount == count, + "Ops.swigluMany: shape mismatch at index \(i)") + precondition( + ups[i].dtype == dtype && outs[i].dtype == dtype, + "Ops.swigluMany: dtype mismatch at index \(i)") + let tgWidth = min(count, 256) + enc.setBuffer(gates[i].buffer, offset: gates[i].offset, index: 0) + enc.setBuffer(ups[i].buffer, offset: ups[i].offset, index: 1) + enc.setBuffer(outs[i].buffer, offset: outs[i].offset, index: 2) + enc.dispatchThreads( + MTLSize(width: count, height: 1, depth: 1), + threadsPerThreadgroup: MTLSize(width: tgWidth, height: 1, depth: 1)) + } + enc.endEncoding() + } + // ─── MoE gather quantised matmul, scalar m1 ──────────────────────── // // Wraps `mt_moe_gather_qmm_int4_{f32,f16,bf16}`. One TG per diff --git a/Sources/FFAI/Ops/OpsFused.swift b/Sources/FFAI/Ops/OpsFused.swift index c6dc7f80..314466f1 100644 --- a/Sources/FFAI/Ops/OpsFused.swift +++ b/Sources/FFAI/Ops/OpsFused.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -31,6 +31,13 @@ // for the attention QKV step at decode (T=1) // * `rmsNormSmall` — RMSNorm specialized for "small" row widths // (n < 128) where the wide variant would be wasteful. +// * `scalarFMA` / `scalarFMAMany` / `scalarFMAChain8` — +// scalar-broadcast FMA for the MoE top-K expert accumulator path +// (avoids materializing a `Tensor.filled([hidden], weight)` +// broadcast buffer per expert). +// * `sigmoidScalarFMAResidual` — `sigmoidScalarFMA` with the +// residual add folded in for the post-MoE-FFN path (the bare +// `sigmoidScalarFMA` sibling lives in `Ops.swift`). import Foundation import Metal @@ -333,4 +340,415 @@ extension Ops { } return result } + + /// `out[i] = base[i] + scalar[0] * value[i]` with `scalar` a + /// 1-element buffer. Replaces the MoE top-K weighted-add chain at + /// decode T=1, which would otherwise be `Tensor.filled([hidden], + /// weight)` + `Ops.mul(expertOut, broadcast)` + `Ops.add(acc, + /// scaled)` — collapsing 1 host alloc + 2 dispatches into 1 + /// dispatch + a 4-byte scalar buffer. Aliasing `out == base` is + /// safe: the kernel reads `value[idx]` and `base[idx]` then writes + /// `out[idx]`. + public static func scalarFMA( + scalar: Tensor, value: Tensor, base: Tensor, + into out: Tensor, on cmd: MTLCommandBuffer + ) { + precondition( + scalar.dtype == value.dtype && value.dtype == base.dtype + && base.dtype == out.dtype, + "Ops.scalarFMA: all tensors must share dtype") + precondition( + scalar.elementCount == 1, + "Ops.scalarFMA: scalar must be [1] (got \(scalar.elementCount))") + precondition( + value.elementCount == base.elementCount + && base.elementCount == out.elementCount, + "Ops.scalarFMA: value / base / out must have matching elementCount") + let n = value.elementCount + let tgWidth = min(n, 256) + let grid = MTLSize(width: n, height: 1, depth: 1) + let tg = MTLSize(width: tgWidth, height: 1, depth: 1) + switch out.dtype { + case .f32: + MetalTileKernels.mt_scalar_fma_f32( + scalar: scalar.buffer, scalarOffset: scalar.offset, + value: value.buffer, valueOffset: value.offset, + base: base.buffer, baseOffset: base.offset, + out: out.buffer, outOffset: out.offset, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.mt_scalar_fma_f16( + scalar: scalar.buffer, scalarOffset: scalar.offset, + value: value.buffer, valueOffset: value.offset, + base: base.buffer, baseOffset: base.offset, + out: out.buffer, outOffset: out.offset, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.mt_scalar_fma_bf16( + scalar: scalar.buffer, scalarOffset: scalar.offset, + value: value.buffer, valueOffset: value.offset, + base: base.buffer, baseOffset: base.offset, + out: out.buffer, outOffset: out.offset, + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError("Ops.scalarFMA: unsupported dtype \(out.dtype)") + } + } + + /// N back-to-back `scalarFMA` dispatches accumulating into the same + /// `acc` tensor on ONE compute encoder. Saves N-1 encoder + /// begin/end pairs versus N independent `scalarFMA` calls. Used by + /// the MoE top-K accumulator path (N=topK, ×nMoELayers per decode + /// token). Metal's in-order execution within a single encoder makes + /// the serial accumulation safe. + public static func scalarFMAMany( + scalars: [Tensor], values: [Tensor], acc: Tensor, + on cmd: MTLCommandBuffer + ) { + precondition( + scalars.count == values.count, "Ops.scalarFMAMany: count mismatch") + precondition(!scalars.isEmpty, "Ops.scalarFMAMany: empty") + let dt = acc.dtype + for (i, s) in scalars.enumerated() { + precondition( + s.elementCount == 1, + "Ops.scalarFMAMany: scalar[\(i)] must be [1]") + precondition( + s.dtype == dt && values[i].dtype == dt, + "Ops.scalarFMAMany: dtype mismatch at \(i)") + precondition( + values[i].elementCount == acc.elementCount, + "Ops.scalarFMAMany: value[\(i)] / acc size mismatch") + } + let n = acc.elementCount + let tgWidth = min(n, 256) + let grid = MTLSize(width: n, height: 1, depth: 1) + let tg = MTLSize(width: tgWidth, height: 1, depth: 1) + guard let enc = cmd.makeComputeCommandEncoder() else { return } + let psoName: String + switch dt { + case .f32: psoName = "mt_scalar_fma_f32" + case .f16: psoName = "mt_scalar_fma_f16" + case .bf16: psoName = "mt_scalar_fma_bf16" + default: fatalError("Ops.scalarFMAMany: unsupported dtype \(dt)") + } + let pso = PSOCache.shared.pipelineState(for: psoName) + enc.setComputePipelineState(pso) + // Kernel buffer bindings: 0 scalar, 1 value, 2 base, 3 out. + // Base and out are always `acc`; only scalar/value rotate per + // dispatch within this encoder. + enc.setBuffer(acc.buffer, offset: acc.offset, index: 2) + enc.setBuffer(acc.buffer, offset: acc.offset, index: 3) + for i in 0 ..< scalars.count { + enc.setBuffer(scalars[i].buffer, offset: scalars[i].offset, index: 0) + enc.setBuffer(values[i].buffer, offset: values[i].offset, index: 1) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + } + enc.endEncoding() + } + + /// 8-way fused scalarFMA chain. Computes + /// `out[i] = sum_{k=0..8} scalars[k][0] * values[k][i]` in ONE + /// kernel dispatch. Collapses the topK=8 expert accumulator chain + /// (8 sequential `mt_scalar_fma` dispatches + 1 zero-fill of `acc`) + /// into a single dispatch with 16 input buffers, saving 7 read + + /// 7 write roundtrips of `acc` per MoE layer. + public static func scalarFMAChain8( + scalars: [Tensor], values: [Tensor], out: Tensor, + on cmd: MTLCommandBuffer + ) { + precondition( + scalars.count == 8 && values.count == 8, + "Ops.scalarFMAChain8: requires exactly 8 of each") + let dt = out.dtype + for i in 0 ..< 8 { + precondition( + scalars[i].elementCount == 1, + "Ops.scalarFMAChain8: scalar[\(i)] must be [1]") + precondition( + scalars[i].dtype == dt && values[i].dtype == dt, + "Ops.scalarFMAChain8: dtype mismatch at \(i)") + precondition( + values[i].elementCount == out.elementCount, + "Ops.scalarFMAChain8: value[\(i)] / out size mismatch") + } + let n = out.elementCount + let tgWidth = min(n, 256) + let grid = MTLSize(width: n, height: 1, depth: 1) + let tg = MTLSize(width: tgWidth, height: 1, depth: 1) + switch dt { + case .f32: + MetalTileKernels.mt_scalar_fma_chain8_f32( + scalar0: scalars[0].buffer, scalar0Offset: scalars[0].offset, + value0: values[0].buffer, value0Offset: values[0].offset, + scalar1: scalars[1].buffer, scalar1Offset: scalars[1].offset, + value1: values[1].buffer, value1Offset: values[1].offset, + scalar2: scalars[2].buffer, scalar2Offset: scalars[2].offset, + value2: values[2].buffer, value2Offset: values[2].offset, + scalar3: scalars[3].buffer, scalar3Offset: scalars[3].offset, + value3: values[3].buffer, value3Offset: values[3].offset, + scalar4: scalars[4].buffer, scalar4Offset: scalars[4].offset, + value4: values[4].buffer, value4Offset: values[4].offset, + scalar5: scalars[5].buffer, scalar5Offset: scalars[5].offset, + value5: values[5].buffer, value5Offset: values[5].offset, + scalar6: scalars[6].buffer, scalar6Offset: scalars[6].offset, + value6: values[6].buffer, value6Offset: values[6].offset, + scalar7: scalars[7].buffer, scalar7Offset: scalars[7].offset, + value7: values[7].buffer, value7Offset: values[7].offset, + out: out.buffer, outOffset: out.offset, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.mt_scalar_fma_chain8_f16( + scalar0: scalars[0].buffer, scalar0Offset: scalars[0].offset, + value0: values[0].buffer, value0Offset: values[0].offset, + scalar1: scalars[1].buffer, scalar1Offset: scalars[1].offset, + value1: values[1].buffer, value1Offset: values[1].offset, + scalar2: scalars[2].buffer, scalar2Offset: scalars[2].offset, + value2: values[2].buffer, value2Offset: values[2].offset, + scalar3: scalars[3].buffer, scalar3Offset: scalars[3].offset, + value3: values[3].buffer, value3Offset: values[3].offset, + scalar4: scalars[4].buffer, scalar4Offset: scalars[4].offset, + value4: values[4].buffer, value4Offset: values[4].offset, + scalar5: scalars[5].buffer, scalar5Offset: scalars[5].offset, + value5: values[5].buffer, value5Offset: values[5].offset, + scalar6: scalars[6].buffer, scalar6Offset: scalars[6].offset, + value6: values[6].buffer, value6Offset: values[6].offset, + scalar7: scalars[7].buffer, scalar7Offset: scalars[7].offset, + value7: values[7].buffer, value7Offset: values[7].offset, + out: out.buffer, outOffset: out.offset, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.mt_scalar_fma_chain8_bf16( + scalar0: scalars[0].buffer, scalar0Offset: scalars[0].offset, + value0: values[0].buffer, value0Offset: values[0].offset, + scalar1: scalars[1].buffer, scalar1Offset: scalars[1].offset, + value1: values[1].buffer, value1Offset: values[1].offset, + scalar2: scalars[2].buffer, scalar2Offset: scalars[2].offset, + value2: values[2].buffer, value2Offset: values[2].offset, + scalar3: scalars[3].buffer, scalar3Offset: scalars[3].offset, + value3: values[3].buffer, value3Offset: values[3].offset, + scalar4: scalars[4].buffer, scalar4Offset: scalars[4].offset, + value4: values[4].buffer, value4Offset: values[4].offset, + scalar5: scalars[5].buffer, scalar5Offset: scalars[5].offset, + value5: values[5].buffer, value5Offset: values[5].offset, + scalar6: scalars[6].buffer, scalar6Offset: scalars[6].offset, + value6: values[6].buffer, value6Offset: values[6].offset, + scalar7: scalars[7].buffer, scalar7Offset: scalars[7].offset, + value7: values[7].buffer, value7Offset: values[7].offset, + out: out.buffer, outOffset: out.offset, + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError("Ops.scalarFMAChain8: unsupported dtype \(dt)") + } + } + + /// Fused MoE phase 1b + phase 2 + phase 3 in ONE kernel launch via + /// `ffai_moe_down_swiglu_accum_int4_chain8`. + /// + /// The GPU-router MoE path runs three back-to-back dispatches per + /// layer: + /// 1. `swigluMany` — `inner[k][i] = silu(gate[k][i]) * up[k][i]` + /// 2. `dequantGemvInt4ExpertIndexedMany` (down) + /// — `out[k] = W_down[expert_idx[k]] · inner[k]` + /// 3. `scalarFMAChain8` — `acc[i] = Σ_k slot_weight[k] · out[k][i]` + /// + /// This wrapper collapses all three into one dispatch. Each + /// threadgroup owns one output row of `[out_dim]`, iterates the 8 + /// slots sequentially: stages `inner` in 3 KiB threadgroup memory, + /// runs the dequant-gemv inner-product against + /// `W_down[expert_idx[k]]` for that slot, accumulates into a + /// per-thread `acc` with the slot scalar baked in, then reduce- + /// sums across threads at the end. Eliminates the `inner[k]` and + /// `out[k]` DRAM roundtrips between phases. + /// + /// Kernel constraints (preconditioned): + /// - `in_dim` (moeIntermediate) MUST be ≤ 768. The kernel's + /// TG-memory alloc is hardcoded at 768 floats (3 KiB). + /// - `group_size` MUST be 64. + /// - TPG = 128. Grid = `[out_dim · TPG, 1, 1]` (one TG per output + /// row). + /// - `gates` / `ups` each contain 8 tensors of shape `[in_dim]` in + /// the model dtype. + /// - `expertIndices` is `[8] u32`. `slotWeights` is `[8] T`. + /// - `weightsStacked` is `[nExperts, out_dim, in_dim/8]` u32-packed. + /// - `scalesStacked` / `biasesStacked` are + /// `[nExperts, out_dim, in_dim/group_size]` T. + /// - `output` is `[out_dim]` T. + public static func moeDownSwigluAccumInt4Chain8( + gates: [Tensor], // 8 tensors, each [moeIntermediate] + ups: [Tensor], // 8 tensors, each [moeIntermediate] + expertIndices: Tensor, // [8] u32 + slotWeights: Tensor, // [8] T + weightsStacked: Tensor, // [nExperts, hidden, moeIntermediate/8] u32 + scalesStacked: Tensor, // [nExperts, hidden, moeIntermediate/groupSize] T + biasesStacked: Tensor, // [nExperts, hidden, moeIntermediate/groupSize] T + output: Tensor, // [hidden] T + inDim: Int, // moeIntermediate + outDim: Int, // hidden + groupSize: Int = 64, + on cmd: MTLCommandBuffer + ) { + precondition( + gates.count == 8 && ups.count == 8, + "Ops.moeDownSwigluAccumInt4Chain8: k must be 8 (got \(gates.count) gates, \(ups.count) ups)") + precondition( + inDim <= 768, + "Ops.moeDownSwigluAccumInt4Chain8: in_dim must be ≤ 768 (kernel TG-mem alloc), got \(inDim)") + precondition( + groupSize == 64, + "Ops.moeDownSwigluAccumInt4Chain8: group_size must be 64, got \(groupSize)") + precondition( + expertIndices.dtype == .u32 && expertIndices.elementCount == 8, + "Ops.moeDownSwigluAccumInt4Chain8: expert_indices must be [8] u32") + precondition( + slotWeights.dtype == output.dtype && slotWeights.elementCount == 8, + "Ops.moeDownSwigluAccumInt4Chain8: slot_weights must be [8] matching output dtype") + precondition( + output.elementCount == outDim, + "Ops.moeDownSwigluAccumInt4Chain8: output must be [out_dim]") + let tpg = 128 + let grid = MTLSize(width: outDim * tpg, height: 1, depth: 1) + let tg = MTLSize(width: tpg, height: 1, depth: 1) + switch output.dtype { + case .f32: + MetalTileKernels.ffai_moe_down_swiglu_accum_int4_chain8_f32( + gate_0: gates[0].buffer, gate_0Offset: gates[0].offset, + up_0: ups[0].buffer, up_0Offset: ups[0].offset, + gate_1: gates[1].buffer, gate_1Offset: gates[1].offset, + up_1: ups[1].buffer, up_1Offset: ups[1].offset, + gate_2: gates[2].buffer, gate_2Offset: gates[2].offset, + up_2: ups[2].buffer, up_2Offset: ups[2].offset, + gate_3: gates[3].buffer, gate_3Offset: gates[3].offset, + up_3: ups[3].buffer, up_3Offset: ups[3].offset, + gate_4: gates[4].buffer, gate_4Offset: gates[4].offset, + up_4: ups[4].buffer, up_4Offset: ups[4].offset, + gate_5: gates[5].buffer, gate_5Offset: gates[5].offset, + up_5: ups[5].buffer, up_5Offset: ups[5].offset, + gate_6: gates[6].buffer, gate_6Offset: gates[6].offset, + up_6: ups[6].buffer, up_6Offset: ups[6].offset, + gate_7: gates[7].buffer, gate_7Offset: gates[7].offset, + up_7: ups[7].buffer, up_7Offset: ups[7].offset, + expert_indices: expertIndices.buffer, expert_indicesOffset: expertIndices.offset, + slot_weights: slotWeights.buffer, slot_weightsOffset: slotWeights.offset, + weights_stacked: weightsStacked.buffer, weights_stackedOffset: weightsStacked.offset, + scales_stacked: scalesStacked.buffer, scales_stackedOffset: scalesStacked.offset, + biases_stacked: biasesStacked.buffer, biases_stackedOffset: biasesStacked.offset, + output: output.buffer, outputOffset: output.offset, + in_dim: UInt32(inDim), out_dim: UInt32(outDim), + group_size: UInt32(groupSize), + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.ffai_moe_down_swiglu_accum_int4_chain8_f16( + gate_0: gates[0].buffer, gate_0Offset: gates[0].offset, + up_0: ups[0].buffer, up_0Offset: ups[0].offset, + gate_1: gates[1].buffer, gate_1Offset: gates[1].offset, + up_1: ups[1].buffer, up_1Offset: ups[1].offset, + gate_2: gates[2].buffer, gate_2Offset: gates[2].offset, + up_2: ups[2].buffer, up_2Offset: ups[2].offset, + gate_3: gates[3].buffer, gate_3Offset: gates[3].offset, + up_3: ups[3].buffer, up_3Offset: ups[3].offset, + gate_4: gates[4].buffer, gate_4Offset: gates[4].offset, + up_4: ups[4].buffer, up_4Offset: ups[4].offset, + gate_5: gates[5].buffer, gate_5Offset: gates[5].offset, + up_5: ups[5].buffer, up_5Offset: ups[5].offset, + gate_6: gates[6].buffer, gate_6Offset: gates[6].offset, + up_6: ups[6].buffer, up_6Offset: ups[6].offset, + gate_7: gates[7].buffer, gate_7Offset: gates[7].offset, + up_7: ups[7].buffer, up_7Offset: ups[7].offset, + expert_indices: expertIndices.buffer, expert_indicesOffset: expertIndices.offset, + slot_weights: slotWeights.buffer, slot_weightsOffset: slotWeights.offset, + weights_stacked: weightsStacked.buffer, weights_stackedOffset: weightsStacked.offset, + scales_stacked: scalesStacked.buffer, scales_stackedOffset: scalesStacked.offset, + biases_stacked: biasesStacked.buffer, biases_stackedOffset: biasesStacked.offset, + output: output.buffer, outputOffset: output.offset, + in_dim: UInt32(inDim), out_dim: UInt32(outDim), + group_size: UInt32(groupSize), + gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.ffai_moe_down_swiglu_accum_int4_chain8_bf16( + gate_0: gates[0].buffer, gate_0Offset: gates[0].offset, + up_0: ups[0].buffer, up_0Offset: ups[0].offset, + gate_1: gates[1].buffer, gate_1Offset: gates[1].offset, + up_1: ups[1].buffer, up_1Offset: ups[1].offset, + gate_2: gates[2].buffer, gate_2Offset: gates[2].offset, + up_2: ups[2].buffer, up_2Offset: ups[2].offset, + gate_3: gates[3].buffer, gate_3Offset: gates[3].offset, + up_3: ups[3].buffer, up_3Offset: ups[3].offset, + gate_4: gates[4].buffer, gate_4Offset: gates[4].offset, + up_4: ups[4].buffer, up_4Offset: ups[4].offset, + gate_5: gates[5].buffer, gate_5Offset: gates[5].offset, + up_5: ups[5].buffer, up_5Offset: ups[5].offset, + gate_6: gates[6].buffer, gate_6Offset: gates[6].offset, + up_6: ups[6].buffer, up_6Offset: ups[6].offset, + gate_7: gates[7].buffer, gate_7Offset: gates[7].offset, + up_7: ups[7].buffer, up_7Offset: ups[7].offset, + expert_indices: expertIndices.buffer, expert_indicesOffset: expertIndices.offset, + slot_weights: slotWeights.buffer, slot_weightsOffset: slotWeights.offset, + weights_stacked: weightsStacked.buffer, weights_stackedOffset: weightsStacked.offset, + scales_stacked: scalesStacked.buffer, scales_stackedOffset: scalesStacked.offset, + biases_stacked: biasesStacked.buffer, biases_stackedOffset: biasesStacked.offset, + output: output.buffer, outputOffset: output.offset, + in_dim: UInt32(inDim), out_dim: UInt32(outDim), + group_size: UInt32(groupSize), + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError("Ops.moeDownSwigluAccumInt4Chain8: unsupported dtype \(output.dtype)") + } + } + + /// `sigmoidScalarFMA` with an extra residual term folded in: + /// `out[i] = residual[i] + base[i] + sigmoid(gate[0]) * value[i]`. + /// Collapses the MoE post-FFN two-step chain + /// `sigmoidScalarFMA(gate, sharedOut, routed) -> ffnOut` followed + /// by `Ops.add(postMix, ffnOut)` into one kernel, saving a + /// `[hidden]` DRAM roundtrip per MoE layer per decode token. + public static func sigmoidScalarFMAResidual( + gate: Tensor, value: Tensor, base: Tensor, residual: Tensor, + into out: Tensor, on cmd: MTLCommandBuffer + ) { + precondition( + gate.dtype == value.dtype && value.dtype == base.dtype + && base.dtype == residual.dtype && residual.dtype == out.dtype, + "Ops.sigmoidScalarFMAResidual: all tensors must share dtype") + precondition( + gate.elementCount == 1, + "Ops.sigmoidScalarFMAResidual: gate must be [1] (got \(gate.elementCount))") + precondition( + value.elementCount == base.elementCount + && base.elementCount == residual.elementCount + && residual.elementCount == out.elementCount, + "Ops.sigmoidScalarFMAResidual: value/base/residual/out must match elementCount") + let n = value.elementCount + let tgWidth = min(n, 256) + let grid = MTLSize(width: n, height: 1, depth: 1) + let tg = MTLSize(width: tgWidth, height: 1, depth: 1) + switch out.dtype { + case .f32: + MetalTileKernels.mt_sigmoid_scalar_fma_residual_f32( + gate: gate.buffer, gateOffset: gate.offset, + value: value.buffer, valueOffset: value.offset, + base: base.buffer, baseOffset: base.offset, + residual: residual.buffer, residualOffset: residual.offset, + out: out.buffer, outOffset: out.offset, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.mt_sigmoid_scalar_fma_residual_f16( + gate: gate.buffer, gateOffset: gate.offset, + value: value.buffer, valueOffset: value.offset, + base: base.buffer, baseOffset: base.offset, + residual: residual.buffer, residualOffset: residual.offset, + out: out.buffer, outOffset: out.offset, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.mt_sigmoid_scalar_fma_residual_bf16( + gate: gate.buffer, gateOffset: gate.offset, + value: value.buffer, valueOffset: value.offset, + base: base.buffer, baseOffset: base.offset, + residual: residual.buffer, residualOffset: residual.offset, + out: out.buffer, outOffset: out.offset, + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError("Ops.sigmoidScalarFMAResidual: unsupported dtype \(out.dtype)") + } + } } diff --git a/Sources/FFAI/Ops/OpsValidation.swift b/Sources/FFAI/Ops/OpsValidation.swift index 22f762b2..b9bfc476 100644 --- a/Sources/FFAI/Ops/OpsValidation.swift +++ b/Sources/FFAI/Ops/OpsValidation.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -173,8 +173,9 @@ public enum OpsValidation { headDim: Int, nQHeads: Int, nKVHeads: Int, baseKV: Int, nQuery: Int, kvStride: Int ) -> String? { - if headDim != 128 { - return "head_dim must be 128 (got \(headDim)); ffai_sdpa_multi is head_dim-128 only" + if headDim != 128 && headDim != 256 { + return "head_dim must be 128 or 256 (got \(headDim)); " + + "ffai_sdpa_multi has dedicated d128 / d256 kernel variants only" } if nQHeads <= 0 { return "nQHeads must be positive (got \(nQHeads))" diff --git a/Sources/MetalTileSwift/MetalTileLibrary.swift b/Sources/MetalTileSwift/MetalTileLibrary.swift index c9af5f59..c4208e6d 100644 --- a/Sources/MetalTileSwift/MetalTileLibrary.swift +++ b/Sources/MetalTileSwift/MetalTileLibrary.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -87,7 +87,13 @@ public final class MetalTileLibrary: @unchecked Sendable { { return parsed } - return 16 + // 64 lets the Metal driver pipeline more cmd-buffers in flight + // before applying backpressure. Measured as a 2-3% prefill + // T=512 win on Qwen3.6-A3B / M5 Max vs the 16 default — the + // bump absorbs per-cmd-buffer encode latency across more + // concurrent submissions. Set FFAI_MAX_COMMAND_BUFFERS to + // override for triage. + return 64 }() /// Process-wide singleton. Lazily initialized; throws on first access if diff --git a/Tests/FFAITests/Benchmark/MoEBgemmBm64MppTests.swift b/Tests/FFAITests/Benchmark/MoEBgemmBm64MppTests.swift index 946aec65..f4250309 100644 --- a/Tests/FFAITests/Benchmark/MoEBgemmBm64MppTests.swift +++ b/Tests/FFAITests/Benchmark/MoEBgemmBm64MppTests.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -34,6 +34,18 @@ import Testing @Suite("MoE bm64_mpp wrapper correctness") struct MoEBgemmBm64MppTests { + /// The `mt_moe_gather_qmm_mma_int4_bm64_mpp_bf16` kernel uses MPP + /// cooperative-tensor (simdgroup_matrix) intrinsics that require + /// Apple GPU Family 9+ (M3 and later). On Family 7 (M1, CI runner) + /// the kernel compiles but produces zeroed output (cos=0.0). Gate + /// the bf16 cells on Family 9 so M1 CI stays green; the kernel is + /// covered by upstream metaltile's own GPU correctness suite on + /// adequate hardware. + static let mppCooperativeTensorAvailable: Bool = { + guard let device = MTLCreateSystemDefaultDevice() else { return false } + return device.supportsFamily(.apple9) + }() + static func pack8(_ nibbles: [UInt32]) -> UInt32 { precondition(nibbles.count == 8) var word: UInt32 = 0 @@ -259,7 +271,9 @@ struct MoEBgemmBm64MppTests { } /// Upstream `..._bf16_matches_m1_clean_tile` shape. - @Test("bm64_mpp bf16 wrapper matches bm16 reference at clean_tile shape") + @Test( + "bm64_mpp bf16 wrapper matches bm16 reference at clean_tile shape", + .enabled(if: mppCooperativeTensorAvailable, "MPP cooperative tensor — Apple GPU Family 9+")) func bf16WrapperCleanTile() { let cosVal = Self.runWrapperCompare( nExperts: 4, tRows: 64, nOut: 64, kIn: 64, groupSize: 32, @@ -272,7 +286,9 @@ struct MoEBgemmBm64MppTests { /// Multi-n-tile + multi-K-block + group_size=64 — matches the bf16 /// multi-tile cell we added upstream that passes at the kernel level. - @Test("bm64_mpp bf16 wrapper matches bm16 reference at multi-tile shape") + @Test( + "bm64_mpp bf16 wrapper matches bm16 reference at multi-tile shape", + .enabled(if: mppCooperativeTensorAvailable, "MPP cooperative tensor — Apple GPU Family 9+")) func bf16WrapperMultiTile() { let cosVal = Self.runWrapperCompare( nExperts: 8, tRows: 128, nOut: 128, kIn: 128, groupSize: 64, @@ -288,7 +304,9 @@ struct MoEBgemmBm64MppTests { /// group_size=64). bf16. Sparse indices — most experts skipped. /// Mirrors the actual production call site that breaks /// `forwardManyEquivalence`. - @Test("bm64_mpp bf16 wrapper matches bm16 reference at qwen36 gate/up shape") + @Test( + "bm64_mpp bf16 wrapper matches bm16 reference at qwen36 gate/up shape", + .enabled(if: mppCooperativeTensorAvailable, "MPP cooperative tensor — Apple GPU Family 9+")) func bf16WrapperQwen36GateUp() { let cosVal = Self.runWrapperCompare( nExperts: 128, tRows: 64, nOut: 768, kIn: 2048, groupSize: 64, @@ -303,7 +321,9 @@ struct MoEBgemmBm64MppTests { /// Qwen3.6-A3B down shape: same nExperts, mTotal, but nOut=2048 /// (32 n-tiles), kIn=768 (24 K-blocks, 12 groups). Down BGEMM is /// the largest n-fanout — most n-tiles per TG dispatch. - @Test("bm64_mpp bf16 wrapper matches bm16 reference at qwen36 down shape") + @Test( + "bm64_mpp bf16 wrapper matches bm16 reference at qwen36 down shape", + .enabled(if: mppCooperativeTensorAvailable, "MPP cooperative tensor — Apple GPU Family 9+")) func bf16WrapperQwen36Down() { let cosVal = Self.runWrapperCompare( nExperts: 128, tRows: 64, nOut: 2048, kIn: 768, groupSize: 64, @@ -321,7 +341,9 @@ struct MoEBgemmBm64MppTests { /// it confirms the bug is in offline `xcrun metal → metallib` compilation /// of MPP cooperative-tensor kernels (likely SDK 26.5 / runtime 26.4.1 /// header drift). - @Test("bm64_mpp bf16 LIVE-COMPILED matches m1 at down shape") + @Test( + "bm64_mpp bf16 LIVE-COMPILED matches m1 at down shape", + .enabled(if: mppCooperativeTensorAvailable, "MPP cooperative tensor — Apple GPU Family 9+")) func bf16LiveCompiledDown() throws { let nExperts = 128 let kIn = 768 @@ -439,7 +461,9 @@ struct MoEBgemmBm64MppTests { /// MPP cooperative tensors specifically, this path will produce the /// correct output (m1-matching) and the wrapper path will produce /// the broken 0.816 cosine output. - @Test("bm64_mpp bf16 raw dispatchThreadgroups matches m1 at down shape") + @Test( + "bm64_mpp bf16 raw dispatchThreadgroups matches m1 at down shape", + .enabled(if: mppCooperativeTensorAvailable, "MPP cooperative tensor — Apple GPU Family 9+")) func bf16RawDispatchThreadgroupsDown() throws { let nExperts = 128 let kIn = 768 @@ -556,7 +580,9 @@ struct MoEBgemmBm64MppTests { /// reproduces the upstream cosine, the bug is input-dependent (only /// triggers on certain bf16 value patterns). If it still drifts at /// ~0.98, the bug is in the Swift dispatch path. - @Test("bm64_mpp bf16 wrapper down shape with sin inputs (upstream match)") + @Test( + "bm64_mpp bf16 wrapper down shape with sin inputs (upstream match)", + .enabled(if: mppCooperativeTensorAvailable, "MPP cooperative tensor — Apple GPU Family 9+")) func bf16WrapperQwen36DownSinInputs() { let nExperts = 128 let kIn = 768 diff --git a/Tests/FFAITests/Generation/DrafterTests.swift b/Tests/FFAITests/Generation/DrafterTests.swift new file mode 100644 index 00000000..75e6dac2 --- /dev/null +++ b/Tests/FFAITests/Generation/DrafterTests.swift @@ -0,0 +1,401 @@ +// Copyright 2026 Tom Turney (@TheTom) +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// Host-side unit test for the n-gram + tree drafters — no GPU needed. + +import Foundation +import Testing + +@testable import FFAI + +@Suite("TreeDrafter — scaffolding") +struct TreeDrafterScaffoldTests { + + @Test("DraftTreeNode size + depth on a linear chain") + func linearChainStructure() { + // Build: 1 → 2 → 3 (depth 3, size 3, single child each). + let leaf = DraftTreeNode(token: 3) + let mid = DraftTreeNode(token: 2, children: [leaf]) + let root = DraftTreeNode(token: 1, children: [mid]) + #expect(root.size == 3) + #expect(root.depth == 3) + #expect(root.children.count == 1) + } + + @Test("DraftTreeNode size + depth on a branching tree") + func branchingTreeStructure() { + // Root → [A, B]; A → [C]; B → [D, E]. size = 6, depth = 3. + let c = DraftTreeNode(token: 3) + let d = DraftTreeNode(token: 4) + let e = DraftTreeNode(token: 5) + let a = DraftTreeNode(token: 1, children: [c]) + let b = DraftTreeNode(token: 2, children: [d, e]) + let root = DraftTreeNode(token: 0, children: [a, b]) + #expect(root.size == 6) + #expect(root.depth == 3) + } + + @Test("Drafter.proposeTreeLinear converts linear propose to a chain tree") + func linearAdapterMatchesChain() { + let the = 1 + let cat = 2 + let sat = 3 + let on = 4 + let history = [the, cat, sat, on, the, cat, sat] + let nGram = NGramDrafter(maxNMatch: 3, minNMatch: 2) + let linear = nGram.propose(history: history, gamma: 2) + guard let tree = nGram.proposeTreeLinear(history: history, maxDepth: 2) + else { + Issue.record("expected a non-nil tree") + return + } + #expect(tree.size == linear.count) + #expect(tree.depth == linear.count) + // Walk single-child chain comparing tokens. + var node: DraftTreeNode? = tree + for tok in linear { + #expect(node?.token == tok) + node = node?.children.first + } + } + + @Test("LinearTreeAdapter wraps any Drafter as a TreeDrafter") + func linearAdapterIsTreeDrafter() { + let drafter = NGramDrafter(maxNMatch: 3, minNMatch: 1) + let tree: TreeDrafter = LinearTreeAdapter(drafter) + let history = [1, 2, 3, 4, 1, 2, 3] + let result = tree.proposeTree( + history: history, maxDepth: 2, maxNodes: 16) + #expect(result != nil) + if let r = result { + var node: DraftTreeNode? = r + while let n = node { + #expect(n.children.count <= 1) + node = n.children.first + } + } + } + + @Test("LinearTreeAdapter returns nil when the underlying drafter is empty") + func linearAdapterEmptyProposal() { + let drafter = NeverDrafter() + let tree = LinearTreeAdapter(drafter) + let result = tree.proposeTree( + history: [1, 2, 3], maxDepth: 4, maxNodes: 16) + #expect(result == nil) + } +} + +@Suite("NGramTreeDrafter — branching tree from n-gram history") +struct NGramTreeDrafterTests { + + @Test("proposeTree branches on multiple continuations of the same n-gram") + func branchesOnMultipleContinuations() { + // History where trigram "the cat sat" appears 3 times with + // different continuations: ["on", "on", "down"]. Top-2 should + // be "on" (count = 2) then "down" (count = 1). + let the = 1 + let cat = 2 + let sat = 3 + let on = 4 + let down = 5 + let pad = 99 + let history = [ + the, cat, sat, on, pad, + the, cat, sat, on, pad, + the, cat, sat, down, pad, + the, cat, sat, + ] + let drafter = NGramTreeDrafter( + maxNMatch: 3, minNMatch: 3, branchingFactor: 2) + guard + let tree = drafter.proposeTree( + history: history, maxDepth: 1, maxNodes: 16) + else { + Issue.record("expected a non-nil tree") + return + } + // Root token = top-1 = "on" (count = 2). + #expect(tree.token == on) + // depth = 1, branching = 2 → root has 2 children: "on" + "down". + #expect(tree.children.count == 2) + #expect(tree.children.map(\.token).sorted() == [on, down].sorted()) + } + + @Test("proposeTree returns nil when the n-gram never repeats") + func nilOnNoMatch() { + let history = [1, 2, 3, 4, 5] + let drafter = NGramTreeDrafter( + maxNMatch: 3, minNMatch: 2, branchingFactor: 2) + let tree = drafter.proposeTree( + history: history, maxDepth: 2, maxNodes: 16) + #expect(tree == nil) + } + + @Test("propose (linear) falls back to a top-1 chain") + func linearFallbackProposeChain() { + // Bigram "a b" → "c" appears 3×, then once more. + let a = 10 + let b = 11 + let c = 12 + let history = [a, b, c, 99, a, b, c, 99, a, b, c, 99, a, b] + let drafter = NGramTreeDrafter( + maxNMatch: 2, minNMatch: 2, branchingFactor: 3) + let linear = drafter.propose(history: history, gamma: 1) + #expect(linear == [c]) + } + + @Test("nodeBudget caps tree growth") + func nodeBudgetCap() { + let t = [1, 2, 3, 4, 5, 1, 2, 3, 4, 6, 1, 2, 3, 4, 7, 1, 2, 3, 4] + let drafter = NGramTreeDrafter( + maxNMatch: 3, minNMatch: 2, branchingFactor: 3) + guard + let tree = drafter.proposeTree( + history: t, maxDepth: 4, maxNodes: 4) + else { + Issue.record("expected non-nil") + return + } + #expect(tree.size <= 4) + } + + @Test("flatten + treeCausalMask: linear chain → lower-triangular mask") + func flattenLinearChain() { + // chain 1 → 2 → 3 → 4 (depth 4, branching 1) + let leaf = DraftTreeNode(token: 4) + let n3 = DraftTreeNode(token: 3, children: [leaf]) + let n2 = DraftTreeNode(token: 2, children: [n3]) + let root = DraftTreeNode(token: 1, children: [n2]) + let (tokens, parent, paths) = root.flatten() + #expect(tokens == [1, 2, 3, 4]) + #expect(parent == [-1, 0, 1, 2]) + #expect(paths[0] == [0]) + #expect(paths[1] == [0, 1]) + #expect(paths[2] == [0, 1, 2]) + #expect(paths[3] == [0, 1, 2, 3]) + // Linear chain mask = lower triangular (each token attends to + // every prior). + let (mask, t) = root.treeCausalMask() + #expect(t == 4) + for i in 0 ..< t { + for j in 0 ..< t { + let m = mask[i * t + j] + let expected: Float = (j <= i) ? 0.0 : -Float.infinity + #expect(m == expected) + } + } + } + + /// Helper: build the branching tree used by verify-case tests. + private func cousinsTree() -> DraftTreeNode { + // 0 (root, tok = 10) + // / \ + // 1 2 (tok = 20, tok = 30) + // /| | + // 3 4 5 (tok = 40, tok = 50, tok = 60) + let n3 = DraftTreeNode(token: 40) + let n4 = DraftTreeNode(token: 50) + let n5 = DraftTreeNode(token: 60) + let n1 = DraftTreeNode(token: 20, children: [n3, n4]) + let n2 = DraftTreeNode(token: 30, children: [n5]) + return DraftTreeNode(token: 10, children: [n1, n2]) + } + + @Test("verify: root mismatch returns empty accepted + the bonus token") + func verifyRootMismatch() { + let tree = cousinsTree() + // Oracle says the target wants token 99 first (not the root's 10). + let oracleHistoryEnd = 99 + let result = tree.verify( + oracleAtHistoryEnd: oracleHistoryEnd, + oracle: { _ in 0 }) // unused — root rejected + #expect(result.acceptedTokens.isEmpty) + #expect(result.bonusToken == 99) + } + + @Test("verify: oracle agrees with path 0→1→3, then diverges") + func verifyAcceptsPathToN3() { + let tree = cousinsTree() + // After history, target wants 10 (root). At root (flat 0), + // target wants 20 (descend into n1). At n1 (flat 1), target + // wants 40 (descend into n3). At n3 (flat 2), target wants + // 999 (off-tree) — stop. + let oracle: (Int) -> Int = { flat in + switch flat { + case 0: return 20 // pick n1 (token = 20) + case 1: return 40 // pick n3 (token = 40) + case 2: return 999 // off-tree, bonus + default: return -1 + } + } + let result = tree.verify(oracleAtHistoryEnd: 10, oracle: oracle) + #expect(result.acceptedTokens == [10, 20, 40]) + #expect(result.bonusToken == 999) + } + + @Test("verify: oracle picks a sibling at depth 1 → only root accepted") + func verifySiblingOnlyRoot() { + let tree = cousinsTree() + let oracle: (Int) -> Int = { flat in + switch flat { + case 0: return 30 // pick n2 over n1 + case 4: return 60 // descend to n5 + case 5: return 999 // off-tree bonus + default: return -1 + } + } + let result = tree.verify(oracleAtHistoryEnd: 10, oracle: oracle) + #expect(result.acceptedTokens == [10, 30, 60]) + #expect(result.bonusToken == 999) + } + + @Test("verify: full path accept on a linear chain tree") + func verifyFullLinearPath() { + let leaf = DraftTreeNode(token: 4) + let n3 = DraftTreeNode(token: 3, children: [leaf]) + let n2 = DraftTreeNode(token: 2, children: [n3]) + let root = DraftTreeNode(token: 1, children: [n2]) + let oracle: (Int) -> Int = { flat in + switch flat { + case 0: return 2 + case 1: return 3 + case 2: return 4 + case 3: return 999 // bonus after the leaf + default: return -1 + } + } + let result = root.verify(oracleAtHistoryEnd: 1, oracle: oracle) + #expect(result.acceptedTokens == [1, 2, 3, 4]) + #expect(result.bonusToken == 999) + } + + @Test("flatten + treeCausalMask: branching tree masks cousins/siblings") + func flattenBranchingTreeMasksCousins() { + // 0 (root) + // / \ + // 1 2 + // /| | + // 3 4 5 + // DFS order: 0, 1, 3, 4, 2, 5 + let n3 = DraftTreeNode(token: 30) + let n4 = DraftTreeNode(token: 40) + let n5 = DraftTreeNode(token: 50) + let n1 = DraftTreeNode(token: 10, children: [n3, n4]) + let n2 = DraftTreeNode(token: 20, children: [n5]) + let root = DraftTreeNode(token: 0, children: [n1, n2]) + let (tokens, parent, _) = root.flatten() + #expect(tokens == [0, 10, 30, 40, 20, 50]) + #expect(parent == [-1, 0, 1, 1, 0, 4]) + + let (mask, t) = root.treeCausalMask() + #expect(t == 6) + // Which flat-indices is each node allowed to attend? + let allow0: Set = [0] + let allow1: Set = [0, 1] + let allow2: Set = [0, 1, 2] + let allow3: Set = [0, 1, 3] // n4 — NOT n3 (sibling) + let allow4: Set = [0, 4] // n2 — NOT n1 / n3 / n4 + let allow5: Set = [0, 4, 5] // n5 — NOT n1 / n3 / n4 + let allowed: [Set] = [ + allow0, allow1, allow2, allow3, allow4, allow5, + ] + + for i in 0 ..< t { + for j in 0 ..< t { + let m = mask[i * t + j] + let expected: Float = + allowed[i].contains(j) ? 0.0 : -Float.infinity + #expect(m == expected) + } + } + } + + @Test("frequency-tie deterministic order") + func frequencyTieDeterministicOrder() { + // Bigram "a b" → "c" once, → "d" once (tie). Top-2 returns + // both; tie-break is by token id ascending → [c, d] when c < d. + let a = 1 + let b = 2 + let c = 3 + let d = 4 + let history = [a, b, c, 99, a, b, d, 99, a, b] + let drafter = NGramTreeDrafter( + maxNMatch: 2, minNMatch: 2, branchingFactor: 2) + guard + let tree = drafter.proposeTree( + history: history, maxDepth: 1, maxNodes: 8) + else { + Issue.record("expected non-nil") + return + } + // Both c and d have count = 1; tie-break ascending → root = c. + #expect(tree.token == c) + #expect(tree.children.map(\.token) == [c, d]) + } +} + +@Suite("Drafter — n-gram prompt-lookup") +struct DrafterTests { + + @Test("NGramDrafter proposes the continuation after a repeated trigram") + func nGramFindsTrigramContinuation() throws { + let the = 1 + let cat = 2 + let sat = 3 + let on = 4 + let mat = 5 + let dog = 6 + let ran = 7 + // [the, cat, sat, on, the, mat, dog, ran, the, cat, sat] + // Latest 3 = "the cat sat"; earliest at indices 0..2; + // next 3 after that = [on, the, mat]. + let history = [the, cat, sat, on, the, mat, dog, ran, the, cat, sat] + + let drafter = NGramDrafter(maxNMatch: 3, minNMatch: 2) + let proposal = drafter.propose(history: history, gamma: 3) + #expect( + proposal == [on, the, mat], + "expected [on, the, mat], got \(proposal)") + } + + @Test("NGramDrafter falls back to a shorter match when longest absent") + func nGramBigramFallback() throws { + let a = 1 + let b = 2 + let c = 3 + let d = 4 + // [a, b, c, d, a, b] — bigram "a, b" repeats; trigram absent. + let history = [a, b, c, d, a, b] + let drafter = NGramDrafter(maxNMatch: 3, minNMatch: 2) + let proposal = drafter.propose(history: history, gamma: 2) + #expect( + proposal == [c, d], + "expected bigram fallback [c, d], got \(proposal)") + } + + @Test("NGramDrafter returns empty when no match is found") + func nGramNoMatchReturnsEmpty() throws { + let history = [1, 2, 3, 4, 5] + let drafter = NGramDrafter(maxNMatch: 3, minNMatch: 2) + let proposal = drafter.propose(history: history, gamma: 3) + #expect(proposal.isEmpty, "expected empty, got \(proposal)") + } + + @Test("NeverDrafter always proposes empty") + func neverDrafterEmpty() throws { + let drafter = NeverDrafter() + #expect(drafter.propose(history: [1, 2, 3], gamma: 4).isEmpty) + } +} diff --git a/Tests/FFAITests/Ops/OpsSpecialPathTests.swift b/Tests/FFAITests/Ops/OpsSpecialPathTests.swift index 3b2dc65f..d4faff49 100644 --- a/Tests/FFAITests/Ops/OpsSpecialPathTests.swift +++ b/Tests/FFAITests/Ops/OpsSpecialPathTests.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -85,6 +85,232 @@ struct OpsSpecialPathTests { } } + @Test("castToF32Two f16 — same as two sequential castToF32 calls") + func castToF32TwoMatchesSequential() { + autoreleasepool { + let n = 4 + let a = Tensor.empty(shape: [n], dtype: .f16) + a.copyIn(from: [Float16(0.5), 2, -2, 4]) + let b = Tensor.empty(shape: [n], dtype: .f16) + b.copyIn(from: [Float16(-1), 0, 1, 0.25]) + let refA = Tensor.empty(shape: [n], dtype: .f32) + let refB = Tensor.empty(shape: [n], dtype: .f32) + runAndWait { cb in + Ops.castToF32(a, into: refA, on: cb) + Ops.castToF32(b, into: refB, on: cb) + } + let outA = Tensor.empty(shape: [n], dtype: .f32) + let outB = Tensor.empty(shape: [n], dtype: .f32) + runAndWait { cb in + Ops.castToF32Two(a, into: outA, b, into: outB, on: cb) + } + let rA = refA.toArray(as: Float.self) + let gA = outA.toArray(as: Float.self) + let rB = refB.toArray(as: Float.self) + let gB = outB.toArray(as: Float.self) + for i in 0 ..< n { + #expect(abs(gA[i] - rA[i]) < 1e-5, "a[\(i)]") + #expect(abs(gB[i] - rB[i]) < 1e-5, "b[\(i)]") + } + } + } + + @Test("castToF32Three bf16 — same as three sequential castToF32 calls") + func castToF32ThreeMatchesSequential() { + autoreleasepool { + // bf16 round-trip — fewer mantissa bits than f16 but still + // exact when source values fit the type. + let n = 4 + let aSrc: [Float] = [0.5, 2, -2, 4] + let bSrc: [Float] = [-1, 0, 1, 0.25] + let cSrc: [Float] = [3.5, -0.125, 1.0, 0] + // bf16 store: pick the top 16 bits of each Float bit pattern. + func toBF16Bits(_ xs: [Float]) -> [UInt16] { + xs.map { UInt16(truncatingIfNeeded: $0.bitPattern >> 16) } + } + let a = Tensor.empty(shape: [n], dtype: .bf16) + a.copyIn(from: toBF16Bits(aSrc)) + let b = Tensor.empty(shape: [n], dtype: .bf16) + b.copyIn(from: toBF16Bits(bSrc)) + let c = Tensor.empty(shape: [n], dtype: .bf16) + c.copyIn(from: toBF16Bits(cSrc)) + let refA = Tensor.empty(shape: [n], dtype: .f32) + let refB = Tensor.empty(shape: [n], dtype: .f32) + let refC = Tensor.empty(shape: [n], dtype: .f32) + runAndWait { cb in + Ops.castToF32(a, into: refA, on: cb) + Ops.castToF32(b, into: refB, on: cb) + Ops.castToF32(c, into: refC, on: cb) + } + let outA = Tensor.empty(shape: [n], dtype: .f32) + let outB = Tensor.empty(shape: [n], dtype: .f32) + let outC = Tensor.empty(shape: [n], dtype: .f32) + runAndWait { cb in + Ops.castToF32Three( + a, into: outA, b, into: outB, c, into: outC, on: cb) + } + let rA = refA.toArray(as: Float.self) + let gA = outA.toArray(as: Float.self) + let rB = refB.toArray(as: Float.self) + let gB = outB.toArray(as: Float.self) + let rC = refC.toArray(as: Float.self) + let gC = outC.toArray(as: Float.self) + for i in 0 ..< n { + #expect(abs(gA[i] - rA[i]) < 1e-5, "a[\(i)]") + #expect(abs(gB[i] - rB[i]) < 1e-5, "b[\(i)]") + #expect(abs(gC[i] - rC[i]) < 1e-5, "c[\(i)]") + } + } + } + + @Test("swigluMany f32 — N=3 separate dispatches share one encoder") + func swigluManyMatchesSequential() { + autoreleasepool { + let n = 8 + // Build three (gate, up) pairs of different sizes — the + // wrapper sets PSO once and dispatches each independently. + let g0 = Tensor.empty(shape: [n], dtype: .f32) + g0.copyIn(from: (0 ..< n).map { Float($0) * 0.1 }) + let u0 = Tensor.empty(shape: [n], dtype: .f32) + u0.copyIn(from: (0 ..< n).map { Float($0 + 1) }) + let g1 = Tensor.empty(shape: [n], dtype: .f32) + g1.copyIn(from: (0 ..< n).map { Float($0) * -0.2 }) + let u1 = Tensor.empty(shape: [n], dtype: .f32) + u1.copyIn(from: (0 ..< n).map { Float($0) * 0.5 + 1 }) + let g2 = Tensor.empty(shape: [n], dtype: .f32) + g2.copyIn(from: (0 ..< n).map { Float($0) * 0.3 - 0.1 }) + let u2 = Tensor.empty(shape: [n], dtype: .f32) + u2.copyIn(from: (0 ..< n).map { Float(n - $0) }) + var ref0: Tensor! + var ref1: Tensor! + var ref2: Tensor! + runAndWait { cb in + ref0 = Ops.swiglu(gate: g0, up: u0, on: cb) + ref1 = Ops.swiglu(gate: g1, up: u1, on: cb) + ref2 = Ops.swiglu(gate: g2, up: u2, on: cb) + } + let out0 = Tensor.empty(shape: [n], dtype: .f32) + let out1 = Tensor.empty(shape: [n], dtype: .f32) + let out2 = Tensor.empty(shape: [n], dtype: .f32) + runAndWait { cb in + Ops.swigluMany( + gates: [g0, g1, g2], + ups: [u0, u1, u2], + outs: [out0, out1, out2], + on: cb) + } + let pairs: [(Tensor, Tensor)] = [(ref0, out0), (ref1, out1), (ref2, out2)] + for (idx, pair) in pairs.enumerated() { + let r = pair.0.toArray(as: Float.self) + let g = pair.1.toArray(as: Float.self) + for i in 0 ..< n { + #expect(abs(r[i] - g[i]) < 1e-5, "pair \(idx) [\(i)]") + } + } + } + } + + @Test("gatedMixerNormMany f32 — T=2 matches two sequential gatedMixerNorm calls") + func gatedMixerNormManyMatchesSequential() { + autoreleasepool { + let t = 2 + let hv = 2 + let dv = 8 + let total = t * hv * dv + let perToken = hv * dv + let yData = (0 ..< total).map { Float($0 + 1) * 0.1 } + let zData = (0 ..< total).map { Float16(Float($0 % 5) * 0.2 - 0.4) } + let weightData = (0 ..< dv).map { Float16(0.5 + Float($0) * 0.05) } + let weight = Tensor.empty(shape: [dv], dtype: .f16) + weight.copyIn(from: weightData) + let epsBuf = Tensor.empty(shape: [1], dtype: .f32) + epsBuf.copyIn(from: [Float(1e-6)]) + // Reference: two single-token gatedMixerNorm dispatches with + // pre-sliced inputs. All allocations + slicing happen up + // front; runAndWait only carries kernel work. + var refSlices: [Tensor] = [] + for tt in 0 ..< t { + let ySlice = Tensor.empty(shape: [hv, dv], dtype: .f32) + ySlice.copyIn(from: Array(yData[tt * perToken ..< (tt + 1) * perToken])) + let zSlice = Tensor.empty(shape: [hv, dv], dtype: .f16) + zSlice.copyIn(from: Array(zData[tt * perToken ..< (tt + 1) * perToken])) + let outSlice = Tensor.empty(shape: [hv, dv], dtype: .f16) + runAndWait { cb in + Ops.gatedMixerNorm( + y: ySlice, z: zSlice, weight: weight, epsBuf: epsBuf, + into: outSlice, + numValueHeads: hv, valueHeadDim: dv, on: cb) + } + refSlices.append(outSlice) + } + // Batched dispatch over the full T·Hv·Dv span. + let y = Tensor.empty(shape: [t, hv, dv], dtype: .f32) + y.copyIn(from: yData) + let z = Tensor.empty(shape: [t, hv, dv], dtype: .f16) + z.copyIn(from: zData) + let outMany = Tensor.empty(shape: [t, hv, dv], dtype: .f16) + runAndWait { cb in + Ops.gatedMixerNormMany( + y: y, z: z, weight: weight, epsBuf: epsBuf, + into: outMany, + t: t, numValueHeads: hv, valueHeadDim: dv, on: cb) + } + let g = outMany.toArray(as: Float16.self) + for tt in 0 ..< t { + let r = refSlices[tt].toArray(as: Float16.self) + for i in 0 ..< perToken { + let rf = Float(r[i]) + let gf = Float(g[tt * perToken + i]) + #expect(abs(rf - gf) < 5e-3, "t=\(tt) i=\(i): ref \(rf) vs got \(gf)") + } + } + } + } + + @Test("siluCastF32PlusCastF32Two f16 — same as sequential silu+two casts") + func siluCastF32PlusCastF32TwoFromF16() { + autoreleasepool { + let n = 4 + // Reference: three separate dispatches. + let s = Tensor.empty(shape: [n], dtype: .f16) + s.copyIn(from: [Float16(0), 1, -1, 2]) + let a = Tensor.empty(shape: [n], dtype: .f16) + a.copyIn(from: [Float16(0.5), 2, -2, 4]) + let b = Tensor.empty(shape: [n], dtype: .f16) + b.copyIn(from: [Float16(-1), 0, 1, 0.25]) + let refSilu = Tensor.empty(shape: [n], dtype: .f32) + let refA = Tensor.empty(shape: [n], dtype: .f32) + let refB = Tensor.empty(shape: [n], dtype: .f32) + runAndWait { cb in + Ops.siluCastToF32(s, into: refSilu, on: cb) + Ops.castToF32(a, into: refA, on: cb) + Ops.castToF32(b, into: refB, on: cb) + } + // Batched: one encoder, switches PSO between silu+cast and plain cast. + let outSilu = Tensor.empty(shape: [n], dtype: .f32) + let outA = Tensor.empty(shape: [n], dtype: .f32) + let outB = Tensor.empty(shape: [n], dtype: .f32) + runAndWait { cb in + Ops.siluCastF32PlusCastF32Two( + siluIn: s, into: outSilu, + a, into: outA, + b, into: outB, + on: cb) + } + let rRef = refSilu.toArray(as: Float.self) + let rGot = outSilu.toArray(as: Float.self) + let aRef = refA.toArray(as: Float.self) + let aGot = outA.toArray(as: Float.self) + let bRef = refB.toArray(as: Float.self) + let bGot = outB.toArray(as: Float.self) + for i in 0 ..< n { + #expect(abs(rGot[i] - rRef[i]) < 1e-5, "silu[\(i)]") + #expect(abs(aGot[i] - aRef[i]) < 1e-5, "a[\(i)]") + #expect(abs(bGot[i] - bRef[i]) < 1e-5, "b[\(i)]") + } + } + } + @Test("swiglu f32 — out[i] = silu(gate[i]) * up[i]") func swigluF32() { autoreleasepool { @@ -128,8 +354,268 @@ struct OpsSpecialPathTests { } } + @Test("sigmoidScalarFMAResidual f32 — out = residual + base + sigmoid(gate) * value") + func sigmoidScalarFMAResidualF32() { + autoreleasepool { + let gate = Tensor.empty(shape: [1], dtype: .f32) + gate.copyIn(from: [Float(0)]) // sigmoid(0) = 0.5 + let value = Tensor.empty(shape: [4], dtype: .f32) + value.copyIn(from: [Float(2), 4, 6, 8]) + let base = Tensor.empty(shape: [4], dtype: .f32) + base.copyIn(from: [Float(10), 20, 30, 40]) + let residual = Tensor.empty(shape: [4], dtype: .f32) + residual.copyIn(from: [Float(100), 200, 300, 400]) + let out = Tensor.empty(shape: [4], dtype: .f32) + out.zero() + runAndWait { cb in + Ops.sigmoidScalarFMAResidual( + gate: gate, value: value, base: base, residual: residual, + into: out, on: cb) + } + // out[i] = residual[i] + base[i] + 0.5 * value[i] + let r = out.toArray(as: Float.self) + #expect(abs(r[0] - 111) < 1e-4) + #expect(abs(r[1] - 222) < 1e-4) + #expect(abs(r[2] - 333) < 1e-4) + #expect(abs(r[3] - 444) < 1e-4) + } + } + + @Test("scalarFMA f32 — out = base + scalar * value") + func scalarFMAF32() { + autoreleasepool { + let scalar = Tensor.empty(shape: [1], dtype: .f32) + scalar.copyIn(from: [Float(0.25)]) + let value = Tensor.empty(shape: [4], dtype: .f32) + value.copyIn(from: [Float(4), 8, 12, 16]) + let base = Tensor.empty(shape: [4], dtype: .f32) + base.copyIn(from: [Float(10), 20, 30, 40]) + let out = Tensor.empty(shape: [4], dtype: .f32) + out.zero() + runAndWait { cb in + Ops.scalarFMA( + scalar: scalar, value: value, base: base, + into: out, on: cb) + } + // 0.25 * [4,8,12,16] = [1,2,3,4]; + base = [11,22,33,44] + let r = out.toArray(as: Float.self) + #expect(abs(r[0] - 11) < 1e-4) + #expect(abs(r[1] - 22) < 1e-4) + #expect(abs(r[2] - 33) < 1e-4) + #expect(abs(r[3] - 44) < 1e-4) + } + } + + @Test("scalarFMAMany f32 — N=3 serial accumulate into acc on one encoder") + func scalarFMAManyF32() { + autoreleasepool { + let n = 4 + // Three (scalar, value) pairs accumulating into acc starting + // from initial values. Final acc[i] = acc0[i] + sum_k s_k * v_k[i]. + let scalars: [Tensor] = (0 ..< 3).map { _ in + Tensor.empty(shape: [1], dtype: .f32) + } + scalars[0].copyIn(from: [Float(1.0)]) + scalars[1].copyIn(from: [Float(2.0)]) + scalars[2].copyIn(from: [Float(0.5)]) + let values: [Tensor] = (0 ..< 3).map { _ in + Tensor.empty(shape: [n], dtype: .f32) + } + values[0].copyIn(from: [Float(1), 2, 3, 4]) + values[1].copyIn(from: [Float(10), 20, 30, 40]) + values[2].copyIn(from: [Float(100), 200, 300, 400]) + let acc = Tensor.empty(shape: [n], dtype: .f32) + acc.copyIn(from: [Float(1000), 1000, 1000, 1000]) + runAndWait { cb in + Ops.scalarFMAMany(scalars: scalars, values: values, acc: acc, on: cb) + } + // acc[i] = 1000 + 1·v0[i] + 2·v1[i] + 0.5·v2[i] + // i=0: 1000 + 1 + 20 + 50 = 1071 + // i=1: 1000 + 2 + 40 + 100 = 1142 + // i=2: 1000 + 3 + 60 + 150 = 1213 + // i=3: 1000 + 4 + 80 + 200 = 1284 + let r = acc.toArray(as: Float.self) + #expect(abs(r[0] - 1071) < 1e-3) + #expect(abs(r[1] - 1142) < 1e-3) + #expect(abs(r[2] - 1213) < 1e-3) + #expect(abs(r[3] - 1284) < 1e-3) + } + } + + @Test("scalarFMAChain8 f32 — out = sum_{k=0..8} scalar_k * value_k") + func scalarFMAChain8F32() { + autoreleasepool { + let n = 4 + let scalars: [Tensor] = (0 ..< 8).map { i in + let t = Tensor.empty(shape: [1], dtype: .f32) + t.copyIn(from: [Float(i + 1)]) // 1, 2, 3, ..., 8 + return t + } + let values: [Tensor] = (0 ..< 8).map { i in + let t = Tensor.empty(shape: [n], dtype: .f32) + t.copyIn(from: (0 ..< n).map { Float((i + 1) * ($0 + 1)) }) + return t + } + let out = Tensor.empty(shape: [n], dtype: .f32) + out.zero() + runAndWait { cb in + Ops.scalarFMAChain8(scalars: scalars, values: values, out: out, on: cb) + } + // value_k[j] = (k+1)·(j+1); scalar_k = k+1 + // out[j] = sum_{k=0..7} (k+1)^2 · (j+1) = (j+1) · sum k^2 (k=1..8) = (j+1) · 204 + let r = out.toArray(as: Float.self) + for j in 0 ..< n { + let expected = Float((j + 1) * 204) + #expect(abs(r[j] - expected) < 1e-3) + } + } + } + + // MARK: - sdpaDecode2Pass + + @Test("sdpaDecode2Pass f32 — blocks=32 matches single-pass sdpaDecode") + func sdpaDecode2PassMatchesSinglePass() { + autoreleasepool { + // Pin the dispatch geometry contract that PR #10 violated: + // the wrapper preconditions blocks ≥ 32 && blocks % 32 == 0 + // (pass-2 hardcodes a 32-lane block-merge). Cross-check that + // at the minimal legal `blocks=32`, the two-pass output + // matches the canonical single-pass `sdpaDecode` element- + // for-element within bf16-tier tolerance. + let nHeads = 4 + let nKVHeads = 2 + let headDim = 128 // a supported single-pass dim + let nKV = 64 + let kvStride = 64 + let blocks = 32 // the minimum legal value + let scale = 1.0 / Float(Double(headDim).squareRoot()) + + let q = Tensor.empty(shape: [nHeads, headDim], dtype: .f32) + q.copyIn(from: (0 ..< nHeads * headDim).map { Float($0 % 17) * 0.1 - 0.5 }) + let k = Tensor.empty(shape: [nKVHeads, kvStride, headDim], dtype: .f32) + k.copyIn( + from: (0 ..< nKVHeads * kvStride * headDim).map { + Float($0 % 13) * 0.05 - 0.3 + }) + let v = Tensor.empty(shape: [nKVHeads, kvStride, headDim], dtype: .f32) + v.copyIn( + from: (0 ..< nKVHeads * kvStride * headDim).map { + Float($0 % 11) * 0.07 - 0.4 + }) + + // Single-pass reference. + let outRef = Tensor.empty(shape: [nHeads, headDim], dtype: .f32) + runAndWait { cb in + _ = Ops.sdpaDecode( + q: q, k: k, v: v, + nQHeads: nHeads, nKVHeads: nKVHeads, headDim: headDim, + nKV: nKV, kvStride: kvStride, + scale: scale, on: cb, into: outRef) + } + + // Two-pass under test. + let partialO = Tensor.empty( + shape: [nHeads, blocks, headDim], dtype: .f32) + let partialM = Tensor.empty(shape: [nHeads, blocks], dtype: .f32) + let partialL = Tensor.empty(shape: [nHeads, blocks], dtype: .f32) + let out = Tensor.empty(shape: [nHeads, headDim], dtype: .f32) + runAndWait { cb in + Ops.sdpaDecode2Pass( + q: q, k: k, v: v, + nQHeads: nHeads, nKVHeads: nKVHeads, headDim: headDim, + nKV: nKV, kvStride: kvStride, blocks: blocks, + scale: scale, + partialO: partialO, partialM: partialM, partialL: partialL, + into: out, on: cb) + } + + let ref = outRef.toArray(as: Float.self) + let got = out.toArray(as: Float.self) + #expect(ref.count == got.count) + for i in 0 ..< ref.count { + #expect( + abs(got[i] - ref[i]) < 1e-4, + "row=\(i / headDim) d=\(i % headDim): two-pass \(got[i]) vs single-pass \(ref[i])" + ) + } + } + } + // MARK: - KV cache append + @Test("kvCacheUpdateKVMany f32 — writes T rows at the right positions") + func kvCacheUpdateKVManyF32() { + autoreleasepool { + let t = 2 + let nKV = 2 + let headDim = 4 + let maxSeq = 4 + let kSrc = Tensor.empty(shape: [t, nKV, headDim], dtype: .f32) + // Token 0: heads [1..4, 5..8]; token 1: heads [9..12, 13..16] + kSrc.copyIn(from: (0 ..< t * nKV * headDim).map { Float($0 + 1) }) + let vSrc = Tensor.empty(shape: [t, nKV, headDim], dtype: .f32) + vSrc.copyIn(from: (0 ..< t * nKV * headDim).map { Float($0 + 1) * 10 }) + let kCache = Tensor.empty(shape: [nKV, maxSeq, headDim], dtype: .f32) + kCache.zero() + let vCache = Tensor.empty(shape: [nKV, maxSeq, headDim], dtype: .f32) + vCache.zero() + let positions = Tensor.empty(shape: [t], dtype: .u32) + positions.copyIn(from: [UInt32(1), UInt32(2)]) + runAndWait { cb in + Ops.kvCacheUpdateKVMany( + kSrc: kSrc, kCache: kCache, + vSrc: vSrc, vCache: vCache, + positions: positions, t: t, + nKVHeads: nKV, headDim: headDim, maxSeq: maxSeq, on: cb) + } + let kGot = kCache.toArray(as: Float.self) + // head 0 row 1 ← token0 head0 = [1,2,3,4] + #expect(Array(kGot[4 ..< 8]) == [1, 2, 3, 4]) + // head 0 row 2 ← token1 head0 = [9,10,11,12] + #expect(Array(kGot[8 ..< 12]) == [9, 10, 11, 12]) + // head 1 row 1 ← token0 head1 = [5,6,7,8] + #expect(Array(kGot[20 ..< 24]) == [5, 6, 7, 8]) + // head 1 row 2 ← token1 head1 = [13,14,15,16] + #expect(Array(kGot[24 ..< 28]) == [13, 14, 15, 16]) + // Untouched slots stay zero. + #expect(Array(kGot[0 ..< 4]) == [0, 0, 0, 0]) + #expect(Array(kGot[12 ..< 16]) == [0, 0, 0, 0]) + let vGot = vCache.toArray(as: Float.self) + #expect(Array(vGot[4 ..< 8]) == [10, 20, 30, 40]) + #expect(Array(vGot[8 ..< 12]) == [90, 100, 110, 120]) + } + } + + @Test("kvCacheUpdateKV f32 — appends K and V rows on one encoder") + func kvCacheUpdateKVF32() { + autoreleasepool { + let nKV = 2 + let headDim = 4 + let maxSeq = 3 + let kSrc = Tensor.empty(shape: [nKV, headDim], dtype: .f32) + kSrc.copyIn(from: [Float(1), 2, 3, 4, 5, 6, 7, 8]) + let vSrc = Tensor.empty(shape: [nKV, headDim], dtype: .f32) + vSrc.copyIn(from: [Float(10), 20, 30, 40, 50, 60, 70, 80]) + let kCache = Tensor.empty(shape: [nKV, maxSeq, headDim], dtype: .f32) + kCache.zero() + let vCache = Tensor.empty(shape: [nKV, maxSeq, headDim], dtype: .f32) + vCache.zero() + runAndWait { cb in + Ops.kvCacheUpdateKV( + kSrc: kSrc, kCache: kCache, + vSrc: vSrc, vCache: vCache, + nKVHeads: nKV, headDim: headDim, + maxSeq: maxSeq, position: 1, on: cb) + } + let kGot = kCache.toArray(as: Float.self) + #expect(Array(kGot[4 ..< 8]) == [1, 2, 3, 4]) + #expect(Array(kGot[16 ..< 20]) == [5, 6, 7, 8]) + let vGot = vCache.toArray(as: Float.self) + #expect(Array(vGot[4 ..< 8]) == [10, 20, 30, 40]) + #expect(Array(vGot[16 ..< 20]) == [50, 60, 70, 80]) + } + } + @Test("kvCacheUpdate f32 — writes one row into [nKV, maxSeq, headDim]") func kvCacheUpdateF32() { autoreleasepool { diff --git a/Tests/FFAITests/Ops/OpsTests.swift b/Tests/FFAITests/Ops/OpsTests.swift index a45466e3..fe5879b7 100644 --- a/Tests/FFAITests/Ops/OpsTests.swift +++ b/Tests/FFAITests/Ops/OpsTests.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -110,6 +110,93 @@ struct OpsTests { } } + @Test("ropePartialTwo f32 — same result as two sequential ropePartial calls") + func ropePartialTwoMatchesSequential() { + autoreleasepool { + let headDim = 4 + let rotaryDim = 2 + let q0 = Tensor.empty(shape: [headDim], dtype: .f32) + q0.copyIn(from: [Float(1), 0, 7, 9]) + let k0 = Tensor.empty(shape: [headDim], dtype: .f32) + k0.copyIn(from: [Float(0), 1, -3, 4]) + let q1 = Tensor.empty(shape: [headDim], dtype: .f32) + q1.copyIn(from: [Float(1), 0, 7, 9]) + let k1 = Tensor.empty(shape: [headDim], dtype: .f32) + k1.copyIn(from: [Float(0), 1, -3, 4]) + // Reference: two independent ropePartial calls. + runAndWait { cb in + Ops.ropePartial( + q0, position: 3, headDim: headDim, rotaryDim: rotaryDim, + thetaBase: 10000, on: cb) + Ops.ropePartial( + k0, position: 3, headDim: headDim, rotaryDim: rotaryDim, + thetaBase: 10000, on: cb) + } + // Batched: one-encoder, two dispatches. + runAndWait { cb in + Ops.ropePartialTwo( + q1, k1, position: 3, headDim: headDim, rotaryDim: rotaryDim, + thetaBase: 10000, on: cb) + } + let r0q = q0.toArray(as: Float.self) + let r0k = k0.toArray(as: Float.self) + let r1q = q1.toArray(as: Float.self) + let r1k = k1.toArray(as: Float.self) + for i in 0 ..< headDim { + #expect(abs(r0q[i] - r1q[i]) < 1e-5, "q[\(i)]: \(r0q[i]) vs \(r1q[i])") + #expect(abs(r0k[i] - r1k[i]) < 1e-5, "k[\(i)]: \(r0k[i]) vs \(r1k[i])") + } + } + } + + @Test("ropePartialMany f32 — T=3 matches three sequential ropePartial calls") + func ropePartialManyMatchesSequential() { + autoreleasepool { + let t = 3 + let nHeads = 2 + let headDim = 4 + let rotaryDim = 2 + let rowStride = nHeads * headDim + let total = t * rowStride + let data: [Float] = (0 ..< total).map { Float($0 + 1) * 0.1 } + // Reference: T sequential single-token rotations against + // T separate buffers (positions 0, 1, 2). + var refRows: [[Float]] = [] + for tt in 0 ..< t { + let slice = Tensor.empty(shape: [nHeads, headDim], dtype: .f32) + slice.copyIn(from: Array(data[tt * rowStride ..< (tt + 1) * rowStride])) + runAndWait { cb in + Ops.ropePartial( + slice, position: tt, headDim: headDim, rotaryDim: rotaryDim, + thetaBase: 10000, on: cb) + } + refRows.append(slice.toArray(as: Float.self)) + } + // Batched: single dispatch with positions [0, 1, 2]. + let qk = Tensor.empty(shape: [t, nHeads, headDim], dtype: .f32) + qk.copyIn(from: data) + let positions = Tensor.empty(shape: [t], dtype: .u32) + positions.copyIn(from: (0 ..< t).map { UInt32($0) }) + runAndWait { cb in + Ops.ropePartialMany( + qk, positions: positions, + t: t, nHeads: nHeads, + headDim: headDim, rotaryDim: rotaryDim, + rowStride: rowStride, + thetaBase: 10000, on: cb) + } + let got = qk.toArray(as: Float.self) + for tt in 0 ..< t { + let row = refRows[tt] + for i in 0 ..< rowStride { + #expect( + abs(got[tt * rowStride + i] - row[i]) < 1e-5, + "t=\(tt) i=\(i): \(got[tt * rowStride + i]) vs \(row[i])") + } + } + } + } + @Test("ropePartial f32 — rotaryDim == headDim matches full rope") func ropePartialFullEqualsRope() { autoreleasepool { @@ -229,6 +316,29 @@ struct OpsTests { } } + @Test("gatherTwo f32 — two ids streams against one table on one encoder") + func gatherTwoF32() { + autoreleasepool { + let table = Tensor.empty(shape: [3, 2], dtype: .f32) + table.copyIn(from: [Float(10), 11, 20, 21, 30, 31]) + let ids1 = Tensor.empty(shape: [2], dtype: .u32) + ids1.copyIn(from: [UInt32(2), 0]) + let ids2 = Tensor.empty(shape: [3], dtype: .u32) + ids2.copyIn(from: [UInt32(1), 1, 2]) + let out1 = Tensor.empty(shape: [2, 2], dtype: .f32) + let out2 = Tensor.empty(shape: [3, 2], dtype: .f32) + runAndWait { cb in + Ops.gatherTwo( + table: table, + ids1: ids1, into: out1, + ids2: ids2, into: out2, + on: cb) + } + #expect(out1.toArray(as: Float.self) == [30, 31, 10, 11]) + #expect(out2.toArray(as: Float.self) == [20, 21, 20, 21, 30, 31]) + } + } + @Test("gemv f32 — out[i] = sum_j W[i,j] * x[j]") func gemvF32() { autoreleasepool { @@ -357,6 +467,67 @@ struct OpsTests { } } + @Test("rmsNormRowsTwo f32 — matches two sequential rmsNormRows calls") + func rmsNormRowsTwoMatchesSequential() { + autoreleasepool { + let n1 = 128 + let n2 = 256 + let rows = 1 + let xs1: [Float] = (0 ..< n1 * rows).map { Float($0 + 1) } + let ws1: [Float] = (0 ..< n1).map { 1.0 + Float($0 % 7) * 0.05 } + let xs2: [Float] = (0 ..< n2 * rows).map { Float($0 + 1) * 0.3 } + let ws2: [Float] = (0 ..< n2).map { 1.0 + Float($0 % 11) * 0.03 } + // Reference: two sequential single-row rmsNormRows calls. + let x1Ref = Tensor.empty(shape: [rows, n1], dtype: .f32) + x1Ref.copyIn(from: xs1) + let w1Ref = Tensor.empty(shape: [n1], dtype: .f32) + w1Ref.copyIn(from: ws1) + let x2Ref = Tensor.empty(shape: [rows, n2], dtype: .f32) + x2Ref.copyIn(from: xs2) + let w2Ref = Tensor.empty(shape: [n2], dtype: .f32) + w2Ref.copyIn(from: ws2) + var ref1: Tensor! + var ref2: Tensor! + runAndWait { cb in + ref1 = Ops.rmsNormRows( + x1Ref, weight: w1Ref, eps: 1e-6, + nRows: rows, rowSize: n1, on: cb) + ref2 = Ops.rmsNormRows( + x2Ref, weight: w2Ref, eps: 1e-6, + nRows: rows, rowSize: n2, on: cb) + } + // Batched: same eps so the wrapper shares one alloc. + let x1 = Tensor.empty(shape: [rows, n1], dtype: .f32) + x1.copyIn(from: xs1) + let w1 = Tensor.empty(shape: [n1], dtype: .f32) + w1.copyIn(from: ws1) + let x2 = Tensor.empty(shape: [rows, n2], dtype: .f32) + x2.copyIn(from: xs2) + let w2 = Tensor.empty(shape: [n2], dtype: .f32) + w2.copyIn(from: ws2) + let out1 = Tensor.empty(shape: [rows, n1], dtype: .f32) + let out2 = Tensor.empty(shape: [rows, n2], dtype: .f32) + runAndWait { cb in + Ops.rmsNormRowsTwo( + x1, weight: w1, eps1: 1e-6, + nRows1: rows, rowSize1: n1, into: out1, + x2, weight: w2, eps2: 1e-6, + nRows2: rows, rowSize2: n2, into: out2, + on: cb) + } + let r1Ref = ref1.toArray(as: Float.self) + let r2Ref = ref2.toArray(as: Float.self) + let r1 = out1.toArray(as: Float.self) + let r2 = out2.toArray(as: Float.self) + for i in 0 ..< n1 { + #expect(abs(r1[i] - r1Ref[i]) < 1e-4, "norm1[\(i)]: \(r1[i]) vs \(r1Ref[i])") + } + for i in 0 ..< n2 { + #expect(abs(r2[i] - r2Ref[i]) < 1e-4, "norm2[\(i)]: \(r2[i]) vs \(r2Ref[i])") + } + } + } + @Test("rmsNorm f32 — wide row (n=5376) routes to mt_rms_norm_wide") func rmsNormWideF32() { autoreleasepool { @@ -939,6 +1110,222 @@ struct OpsTests { } } + @Test("sdpaDecode2Pass f32 — blocks=32 matches single-pass sdpaDecode") + func sdpaDecode2PassMatchesSinglePass() { + autoreleasepool { + // pass2 hardcodes a 32-lane block reduction, so blocks + // must be ≥ 32 and a multiple of 32 — anything smaller + // silently emits zero output. + let D = 128 + let blocks = 32 + let kvStride = 64 + let nKV = 64 + let nQHeads = 2 + let nKVHeads = 1 + let q = Tensor.empty(shape: [nQHeads, D], dtype: .f32) + var qData = [Float](repeating: 0, count: nQHeads * D) + for h in 0 ..< nQHeads { + qData[h * D] = Float(h + 1) + qData[h * D + 1] = 0.5 + } + q.copyIn(from: qData) + let k = Tensor.empty(shape: [nKVHeads, kvStride, D], dtype: .f32) + var kData = [Float](repeating: 0, count: nKVHeads * kvStride * D) + for p in 0 ..< nKV { + kData[p * D] = Float(p + 1) * 0.05 + kData[p * D + 1] = 0.1 + } + k.copyIn(from: kData) + let v = Tensor.empty(shape: [nKVHeads, kvStride, D], dtype: .f32) + var vData = [Float](repeating: 0, count: nKVHeads * kvStride * D) + for p in 0 ..< nKV { + for d in 0 ..< D { + vData[p * D + d] = Float((p + 1) * (d + 1)) * 0.01 + } + } + v.copyIn(from: vData) + // Reference. + var refOut: Tensor! + runAndWait { cb in + refOut = Ops.sdpaDecode( + q: q, k: k, v: v, + nQHeads: nQHeads, nKVHeads: nKVHeads, + headDim: D, nKV: nKV, kvStride: kvStride, + scale: 1.0, on: cb) + } + // 2-pass. + let partialO = Tensor.empty(shape: [nQHeads, blocks, D], dtype: .f32) + let partialM = Tensor.empty(shape: [nQHeads, blocks], dtype: .f32) + let partialL = Tensor.empty(shape: [nQHeads, blocks], dtype: .f32) + let out = Tensor.empty(shape: [nQHeads, D], dtype: .f32) + partialO.zero() + partialM.zero() + partialL.zero() + out.zero() + runAndWait { cb in + Ops.sdpaDecode2Pass( + q: q, k: k, v: v, + nQHeads: nQHeads, nKVHeads: nKVHeads, headDim: D, + nKV: nKV, kvStride: kvStride, blocks: blocks, + scale: 1.0, + partialO: partialO, partialM: partialM, partialL: partialL, + into: out, on: cb) + } + let r = out.toArray(as: Float.self) + let ref = refOut.toArray(as: Float.self) + for i in 0 ..< nQHeads * D { + #expect(abs(r[i] - ref[i]) < 1e-2, "i=\(i): \(r[i]) vs \(ref[i])") + } + } + } + + @Test("sdpaMultiTreeMask — all-zero mask reproduces non-causal sdpaMulti output") + func sdpaMultiTreeMaskAllowAll() { + autoreleasepool { + // An all-zero additive mask permits every position → + // attention behaves exactly like causal=false sdpaMulti. + let headDim = 128 + let nQHeads = 2 + let nKVHeads = 1 + let baseKV = 0 + let nQuery = 4 + let kvStride = baseKV + nQuery + let scale = 1.0 / Float(Double(headDim).squareRoot()) + let q = Tensor.empty(shape: [nQuery, nQHeads, headDim], dtype: .f32) + q.copyIn(from: (0 ..< nQuery * nQHeads * headDim).map { Float($0 % 7) * 0.1 }) + let k = Tensor.empty(shape: [nKVHeads, kvStride, headDim], dtype: .f32) + k.copyIn(from: [Float](repeating: 0.5, count: nKVHeads * kvStride * headDim)) + var vData = [Float](repeating: 0, count: nKVHeads * kvStride * headDim) + for t in 0 ..< kvStride { + for d in 0 ..< headDim { vData[t * headDim + d] = Float(t) } + } + let v = Tensor.empty(shape: [nKVHeads, kvStride, headDim], dtype: .f32) + v.copyIn(from: vData) + // Reference: plain non-causal sdpaMulti. + var refOut: Tensor! + runAndWait { cb in + refOut = Ops.sdpaMulti( + q: q, k: k, v: v, + nQHeads: nQHeads, nKVHeads: nKVHeads, headDim: headDim, + baseKV: baseKV, nQuery: nQuery, kvStride: kvStride, + causal: false, scale: scale, on: cb) + } + // Tree-mask path: all-zero mask = permit everywhere. + let mask = Tensor.empty(shape: [nQuery, nQuery], dtype: .f32) + mask.copyIn(from: [Float](repeating: 0, count: nQuery * nQuery)) + var got: Tensor! + runAndWait { cb in + got = Ops.sdpaMultiTreeMask( + q: q, k: k, v: v, mask: mask, + nQHeads: nQHeads, nKVHeads: nKVHeads, headDim: headDim, + baseKV: baseKV, nQuery: nQuery, kvStride: kvStride, + scale: scale, on: cb) + } + let r = refOut.toArray(as: Float.self) + let g = got.toArray(as: Float.self) + for i in 0 ..< r.count { + #expect(abs(r[i] - g[i]) < 1e-3, "i=\(i): \(g[i]) vs \(r[i])") + } + } + } + + @Test("sdpaMultiTreeMask — diagonal -inf mask isolates each query to its own row") + func sdpaMultiTreeMaskDiagonalOnly() { + autoreleasepool { + // Mask = -inf off the diagonal, 0 on the diagonal → each + // query attends only its own KV slot (assuming the cache + // is full and queries are appended at the end). Outputs + // collapse to a copy of the matching V row. + let headDim = 128 + let nQHeads = 1 + let nKVHeads = 1 + let baseKV = 0 + let nQuery = 4 + let kvStride = baseKV + nQuery + let scale = 1.0 / Float(Double(headDim).squareRoot()) + // Make every Q row identical and every K row identical so + // attention only depends on the mask, not the dot products. + let q = Tensor.empty(shape: [nQuery, nQHeads, headDim], dtype: .f32) + q.copyIn(from: [Float](repeating: 0.5, count: nQuery * nQHeads * headDim)) + let k = Tensor.empty(shape: [nKVHeads, kvStride, headDim], dtype: .f32) + k.copyIn(from: [Float](repeating: 0.5, count: nKVHeads * kvStride * headDim)) + // V row t holds the constant value `t`. + var vData = [Float](repeating: 0, count: nKVHeads * kvStride * headDim) + for t in 0 ..< kvStride { + for d in 0 ..< headDim { vData[t * headDim + d] = Float(t + 1) } + } + let v = Tensor.empty(shape: [nKVHeads, kvStride, headDim], dtype: .f32) + v.copyIn(from: vData) + let neg: Float = -1e9 + var maskData = [Float](repeating: neg, count: nQuery * nQuery) + for i in 0 ..< nQuery { maskData[i * nQuery + i] = 0 } + let mask = Tensor.empty(shape: [nQuery, nQuery], dtype: .f32) + mask.copyIn(from: maskData) + var got: Tensor! + runAndWait { cb in + got = Ops.sdpaMultiTreeMask( + q: q, k: k, v: v, mask: mask, + nQHeads: nQHeads, nKVHeads: nKVHeads, headDim: headDim, + baseKV: baseKV, nQuery: nQuery, kvStride: kvStride, + scale: scale, on: cb) + } + let g = got.toArray(as: Float.self) + for qIdx in 0 ..< nQuery { + let expected = Float(qIdx + 1) + for d in 0 ..< headDim { + #expect( + abs(g[qIdx * headDim + d] - expected) < 1e-3, + "q=\(qIdx) d=\(d): \(g[qIdx * headDim + d]) vs \(expected)") + } + } + } + } + + @Test("sdpaMulti — head_dim 256 routes to d256 kernel (uniform K → mean V)") + func sdpaMultiHeadDim256UniformKMeansV() { + autoreleasepool { + // Identical to the d=128 uniform-K test but at head_dim=256 + // (Qwen3.6-A3B full-attention layers). Verifies the d256 + // routing exists and the kernel produces the same uniform- + // attention result as the d128 path. + let headDim = 256 + let nQHeads = 2 + let nKVHeads = 1 + let baseKV = 0 + let nQuery = 4 + let kvStride = baseKV + nQuery + let scale = 1.0 / Float(Double(headDim).squareRoot()) + + let q = Tensor.empty(shape: [nQuery, nQHeads, headDim], dtype: .f32) + q.copyIn(from: (0 ..< nQuery * nQHeads * headDim).map { Float($0 % 7) * 0.1 }) + + let k = Tensor.empty(shape: [nKVHeads, kvStride, headDim], dtype: .f32) + k.copyIn(from: [Float](repeating: 0.5, count: nKVHeads * kvStride * headDim)) + + var vData = [Float](repeating: 0, count: nKVHeads * kvStride * headDim) + for t in 0 ..< kvStride { + for d in 0 ..< headDim { vData[t * headDim + d] = Float(t) } + } + let v = Tensor.empty(shape: [nKVHeads, kvStride, headDim], dtype: .f32) + v.copyIn(from: vData) + + let cmd = Device.shared.makeCommandBuffer() + let out = Ops.sdpaMulti( + q: q, k: k, v: v, + nQHeads: nQHeads, nKVHeads: nKVHeads, headDim: headDim, + baseKV: baseKV, nQuery: nQuery, kvStride: kvStride, + causal: false, scale: scale, on: cmd) + cmd.commit() + cmd.waitUntilCompleted() + + let result = out.toArray(as: Float.self) + #expect(result.count == nQuery * nQHeads * headDim) + for value in result { + #expect(abs(value - 1.5) < 1e-3, "expected mean V = 1.5, got \(value)") + } + } + } + @Test("sdpaMulti — causal mode: query r attends V rows 0...r") func sdpaMultiCausalPrefixMeans() { autoreleasepool { diff --git a/Tests/FFAITests/Ops/OpsValidationTests.swift b/Tests/FFAITests/Ops/OpsValidationTests.swift index 2f894063..7b6c19c1 100644 --- a/Tests/FFAITests/Ops/OpsValidationTests.swift +++ b/Tests/FFAITests/Ops/OpsValidationTests.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -255,17 +255,25 @@ struct OpsValidationTests { baseKV: 8, nQuery: 8, kvStride: 4096) == nil) } - @Test("sdpaMulti rejects head_dim != 128") + @Test("sdpaMulti rejects head_dim outside {128, 256}") func sdpaMultiRejectsBadHeadDim() { - // The kernel hardcodes 4 elements/lane (128/32); other head - // dims OOB-read. - for hd in [32, 64, 96, 256] { + // d=128 hardcodes 4 elements/lane (128/32); d=256 has its own + // d256 kernel variant with a 2-phase output reduction. Anything + // else OOB-reads. + for hd in [32, 64, 96, 192, 384, 512] { #expect( OpsValidation.validateSdpaMulti( headDim: hd, nQHeads: 8, nKVHeads: 8, baseKV: 0, nQuery: 8, kvStride: 8) != nil, "head_dim \(hd) should be rejected") } + // headDim=256 must now be ACCEPTED by the validator — d256 + // routing was missing on the dev branch and is restored here. + #expect( + OpsValidation.validateSdpaMulti( + headDim: 256, nQHeads: 8, nKVHeads: 8, + baseKV: 0, nQuery: 8, kvStride: 8) == nil, + "head_dim 256 must be accepted (d256 kernel exists)") } @Test("sdpaMulti rejects non-integer GQA fan-out") diff --git a/Tests/FFAITests/Ops/QuantizedOpsTests.swift b/Tests/FFAITests/Ops/QuantizedOpsTests.swift index a2b2e2bd..ae99ea92 100644 --- a/Tests/FFAITests/Ops/QuantizedOpsTests.swift +++ b/Tests/FFAITests/Ops/QuantizedOpsTests.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -702,4 +702,309 @@ struct QuantizedOpsTests { } } } + + // MARK: - Batched dequantGemvInt4 (Two/Three/Four-projection variants) + // + // The batched wrappers do N back-to-back `dequant_gemv_int4_*` + // dispatches inside one compute encoder, sharing the input + // binding and the `in_dim` / `group_size` constants. Correctness + // = N independent `dequantGemvInt4` calls produce the same output + // tensors. The tests below build one input + N (weight, scales, + // biases) triples, run both code paths, and require element-wise + // equality. + + /// Build one realistic (weight, scales, biases, expectedOut) shape + /// for a single int4 projection. Returns CPU-truth output for + /// later equality checking. + private static func makeInt4Projection( + outDim: Int, inDim: Int, gs: Int, + rowSeed: Int, scaleStep: Float, biasStep: Float, + input: [Float] + ) -> (weight: Tensor, scales: Tensor, biases: Tensor, expected: [Float]) { + let nGroups = inDim / gs + var q = [[UInt32]](repeating: [], count: outDim) + for r in 0 ..< outDim { + q[r] = (0 ..< inDim).map { UInt32(($0 + r * rowSeed) % 16) } + } + let scales: [Float] = (0 ..< (outDim * nGroups)).map { Float($0 + 1) * scaleStep } + let biases: [Float] = (0 ..< (outDim * nGroups)).map { Float($0) * biasStep } + var packed: [UInt32] = [] + for r in 0 ..< outDim { + for i in stride(from: 0, to: inDim, by: 8) { + let nibbles = Array(q[r][i ..< i + 8]) + packed.append(Self.pack8(nibbles)) + } + } + var expected: [Float] = [] + for r in 0 ..< outDim { + var acc: Float = 0 + for g in 0 ..< nGroups { + let s = scales[r * nGroups + g] + let b = biases[r * nGroups + g] + for j in 0 ..< gs { + let qv = Float(q[r][g * gs + j]) + acc += (qv * s + b) * input[g * gs + j] + } + } + expected.append(acc) + } + let weight = Tensor.empty(shape: [outDim, inDim / 8], dtype: .u32) + weight.copyIn(from: packed) + let scalesT = Tensor.empty(shape: [outDim, nGroups], dtype: .f32) + scalesT.copyIn(from: scales) + let biasesT = Tensor.empty(shape: [outDim, nGroups], dtype: .f32) + biasesT.copyIn(from: biases) + return (weight, scalesT, biasesT, expected) + } + + @Test("dequantGemvInt4Two: two projections in one encoder match CPU") + func dequantGemvInt4TwoCorrectness() { + autoreleasepool { + let outDim = 4 + let inDim = 128 + let gs = 64 + let input: [Float] = (0 ..< inDim).map { Float($0) * 0.1 - 6.4 } + let p0 = Self.makeInt4Projection( + outDim: outDim, inDim: inDim, gs: gs, + rowSeed: 7, scaleStep: 0.01, biasStep: -0.005, input: input) + let p1 = Self.makeInt4Projection( + outDim: outDim, inDim: inDim, gs: gs, + rowSeed: 13, scaleStep: 0.02, biasStep: 0.003, input: input) + let inputT = Tensor.empty(shape: [inDim], dtype: .f32) + inputT.copyIn(from: input) + let out0 = Tensor.empty(shape: [outDim], dtype: .f32) + let out1 = Tensor.empty(shape: [outDim], dtype: .f32) + runAndWait { cb in + Ops.dequantGemvInt4Two( + input: inputT, + w0: p0.weight, s0: p0.scales, b0: p0.biases, out0: out0, + w1: p1.weight, s1: p1.scales, b1: p1.biases, out1: out1, + groupSize: gs, on: cb) + } + let r0 = out0.toArray(as: Float.self) + let r1 = out1.toArray(as: Float.self) + for i in 0 ..< outDim { + #expect(abs(r0[i] - p0.expected[i]) < 1e-2) + #expect(abs(r1[i] - p1.expected[i]) < 1e-2) + } + } + } + + @Test("dequantGemvInt4Three: three projections in one encoder match CPU") + func dequantGemvInt4ThreeCorrectness() { + autoreleasepool { + let outDim = 4 + let inDim = 128 + let gs = 64 + let input: [Float] = (0 ..< inDim).map { Float($0) * 0.05 - 3.2 } + let p0 = Self.makeInt4Projection( + outDim: outDim, inDim: inDim, gs: gs, + rowSeed: 7, scaleStep: 0.01, biasStep: -0.005, input: input) + let p1 = Self.makeInt4Projection( + outDim: outDim, inDim: inDim, gs: gs, + rowSeed: 11, scaleStep: 0.02, biasStep: 0.003, input: input) + let p2 = Self.makeInt4Projection( + outDim: outDim, inDim: inDim, gs: gs, + rowSeed: 17, scaleStep: 0.015, biasStep: -0.002, input: input) + let inputT = Tensor.empty(shape: [inDim], dtype: .f32) + inputT.copyIn(from: input) + let out0 = Tensor.empty(shape: [outDim], dtype: .f32) + let out1 = Tensor.empty(shape: [outDim], dtype: .f32) + let out2 = Tensor.empty(shape: [outDim], dtype: .f32) + runAndWait { cb in + Ops.dequantGemvInt4Three( + input: inputT, + w0: p0.weight, s0: p0.scales, b0: p0.biases, out0: out0, + w1: p1.weight, s1: p1.scales, b1: p1.biases, out1: out1, + w2: p2.weight, s2: p2.scales, b2: p2.biases, out2: out2, + groupSize: gs, on: cb) + } + for i in 0 ..< outDim { + #expect(abs(out0.toArray(as: Float.self)[i] - p0.expected[i]) < 1e-2) + #expect(abs(out1.toArray(as: Float.self)[i] - p1.expected[i]) < 1e-2) + #expect(abs(out2.toArray(as: Float.self)[i] - p2.expected[i]) < 1e-2) + } + } + } + + @Test("dequantGemvInt4Four: four projections in one encoder match CPU") + func dequantGemvInt4FourCorrectness() { + autoreleasepool { + let outDim = 4 + let inDim = 128 + let gs = 64 + let input: [Float] = (0 ..< inDim).map { Float($0) * 0.04 - 2.5 } + let p0 = Self.makeInt4Projection( + outDim: outDim, inDim: inDim, gs: gs, + rowSeed: 7, scaleStep: 0.01, biasStep: -0.005, input: input) + let p1 = Self.makeInt4Projection( + outDim: outDim, inDim: inDim, gs: gs, + rowSeed: 11, scaleStep: 0.02, biasStep: 0.003, input: input) + let p2 = Self.makeInt4Projection( + outDim: outDim, inDim: inDim, gs: gs, + rowSeed: 17, scaleStep: 0.015, biasStep: -0.002, input: input) + let p3 = Self.makeInt4Projection( + outDim: outDim, inDim: inDim, gs: gs, + rowSeed: 23, scaleStep: 0.025, biasStep: 0.001, input: input) + let inputT = Tensor.empty(shape: [inDim], dtype: .f32) + inputT.copyIn(from: input) + let out0 = Tensor.empty(shape: [outDim], dtype: .f32) + let out1 = Tensor.empty(shape: [outDim], dtype: .f32) + let out2 = Tensor.empty(shape: [outDim], dtype: .f32) + let out3 = Tensor.empty(shape: [outDim], dtype: .f32) + runAndWait { cb in + Ops.dequantGemvInt4Four( + input: inputT, + w0: p0.weight, s0: p0.scales, b0: p0.biases, out0: out0, + w1: p1.weight, s1: p1.scales, b1: p1.biases, out1: out1, + w2: p2.weight, s2: p2.scales, b2: p2.biases, out2: out2, + w3: p3.weight, s3: p3.scales, b3: p3.biases, out3: out3, + groupSize: gs, on: cb) + } + for i in 0 ..< outDim { + #expect(abs(out0.toArray(as: Float.self)[i] - p0.expected[i]) < 1e-2) + #expect(abs(out1.toArray(as: Float.self)[i] - p1.expected[i]) < 1e-2) + #expect(abs(out2.toArray(as: Float.self)[i] - p2.expected[i]) < 1e-2) + #expect(abs(out3.toArray(as: Float.self)[i] - p3.expected[i]) < 1e-2) + } + } + } + + @Test("dequantGemvInt4Many: N projections with distinct inputs match standalone calls") + func dequantGemvInt4ManyCorrectness() { + autoreleasepool { + // N=3 projections, each with its own input + output. + let outDim = 4 + let inDim = 128 + let gs = 64 + let inputs: [[Float]] = [ + (0 ..< inDim).map { Float($0) * 0.05 - 3.2 }, + (0 ..< inDim).map { Float($0) * 0.04 - 2.5 }, + (0 ..< inDim).map { Float($0) * 0.03 - 1.8 }, + ] + var experts: [(weight: Tensor, scales: Tensor, biases: Tensor, expected: [Float])] = [] + for i in 0 ..< 3 { + let p = Self.makeInt4Projection( + outDim: outDim, inDim: inDim, gs: gs, + rowSeed: 7 + i * 5, scaleStep: 0.01 + Float(i) * 0.005, + biasStep: -0.005 + Float(i) * 0.002, input: inputs[i]) + experts.append(p) + } + var inputTs: [Tensor] = [] + for i in 0 ..< 3 { + let t = Tensor.empty(shape: [inDim], dtype: .f32) + t.copyIn(from: inputs[i]) + inputTs.append(t) + } + let outs = (0 ..< 3).map { _ in Tensor.empty(shape: [outDim], dtype: .f32) } + runAndWait { cb in + Ops.dequantGemvInt4Many( + weights: experts.map { $0.weight }, + scales: experts.map { $0.scales }, + biases: experts.map { $0.biases }, + inputs: inputTs, outputs: outs, + groupSize: gs, on: cb) + } + for i in 0 ..< 3 { + let got = outs[i].toArray(as: Float.self) + for j in 0 ..< outDim { + #expect(abs(got[j] - experts[i].expected[j]) < 1e-2) + } + } + } + } + + @Test("dequantGemvInt4ExpertIndexed: matches standalone dequantGemvInt4 for the picked expert") + func dequantGemvInt4ExpertIndexedCorrectness() { + autoreleasepool { + let nExperts = 4 + let outDim = 4 + let inDim = 128 + let gs = 64 + let input: [Float] = (0 ..< inDim).map { Float($0) * 0.05 - 3.2 } + // Build per-expert weight tensors AND a stacked weight tensor + // that concatenates each expert's slab. + var experts: [(weight: Tensor, scales: Tensor, biases: Tensor, expected: [Float])] = [] + for e in 0 ..< nExperts { + let p = Self.makeInt4Projection( + outDim: outDim, inDim: inDim, gs: gs, + rowSeed: 7 + e * 5, scaleStep: 0.01 + Float(e) * 0.005, + biasStep: -0.005 + Float(e) * 0.002, input: input) + experts.append(p) + } + // Stack: [nExperts, outDim, inDim/8] + let packedPerRow = inDim / 8 + let nGroups = inDim / gs + let weightsStacked = Tensor.empty( + shape: [nExperts, outDim, packedPerRow], dtype: .u32) + let scalesStacked = Tensor.empty( + shape: [nExperts, outDim, nGroups], dtype: .f32) + let biasesStacked = Tensor.empty( + shape: [nExperts, outDim, nGroups], dtype: .f32) + // Repack by copying each expert's data into the stacked tensor. + var packedAll: [UInt32] = [] + var scalesAll: [Float] = [] + var biasesAll: [Float] = [] + for e in 0 ..< nExperts { + packedAll.append(contentsOf: experts[e].weight.toArray(as: UInt32.self)) + scalesAll.append(contentsOf: experts[e].scales.toArray(as: Float.self)) + biasesAll.append(contentsOf: experts[e].biases.toArray(as: Float.self)) + } + weightsStacked.copyIn(from: packedAll) + scalesStacked.copyIn(from: scalesAll) + biasesStacked.copyIn(from: biasesAll) + + let inputT = Tensor.empty(shape: [inDim], dtype: .f32) + inputT.copyIn(from: input) + + // Pick expert 2. + let pick: UInt32 = 2 + let pickT = Tensor.empty(shape: [1], dtype: .u32) + pickT.copyIn(from: [pick]) + let out = Tensor.empty(shape: [outDim], dtype: .f32) + runAndWait { cb in + Ops.dequantGemvInt4ExpertIndexed( + weightsStacked: weightsStacked, + scalesStacked: scalesStacked, + biasesStacked: biasesStacked, + input: inputT, expertIndex: pickT, + groupSize: gs, on: cb, into: out) + } + let got = out.toArray(as: Float.self) + let want = experts[Int(pick)].expected + for i in 0 ..< outDim { + #expect(abs(got[i] - want[i]) < 1e-2, "row \(i)") + } + } + } + + @Test("moeRouterTopK f32 — top-K of (1, 5, 3, 4) at k=2 picks indices [1, 3]") + func moeRouterTopKDeterministic() { + autoreleasepool { + // norm_topk_prob = true → weights are softmax over the two + // selected logits and renormalise to 1.0. With logits + // (1, 5, 3, 4), the top-2 are indices 1 and 3 (values 5, 4). + // softmax([5, 4]) = (e^5, e^4) / (e^5 + e^4) ≈ (0.731, 0.269). + let nExperts = 4 + let k = 2 + let logits = Tensor.empty(shape: [nExperts], dtype: .f32) + logits.copyIn(from: [Float(1), 5, 3, 4]) + let indicesOut = Tensor.empty(shape: [k], dtype: .u32) + let weightsOut = Tensor.empty(shape: [k], dtype: .f32) + runAndWait { cb in + Ops.moeRouterTopK( + logits: logits, + indicesOut: indicesOut, weightsOut: weightsOut, + nExperts: nExperts, k: k, normTopkProb: true, on: cb) + } + let idx = indicesOut.toArray(as: UInt32.self) + let wts = weightsOut.toArray(as: Float.self) + #expect(Set(idx) == Set([UInt32(1), 3])) + // Sum should be ~1.0 (norm_topk_prob). + #expect(abs(wts[0] + wts[1] - 1.0) < 1e-3) + // Index 1 has the larger logit so its weight should dominate. + let weightOf1 = idx[0] == 1 ? wts[0] : wts[1] + #expect(weightOf1 > 0.7) + } + } } diff --git a/Tests/Helpers/CommonTestHelpers.swift b/Tests/Helpers/CommonTestHelpers.swift index e3daaf41..47c574ba 100644 --- a/Tests/Helpers/CommonTestHelpers.swift +++ b/Tests/Helpers/CommonTestHelpers.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -20,6 +20,23 @@ import FFAI import Foundation +import Metal + +// MARK: - MTLCommandBuffer async-friendly wait + +/// SDK-portable replacement for `await cmd.completed()` (Xcode 26+ +/// only) and `cmd.waitUntilCompleted()` (now flagged unavailable from +/// async contexts in Xcode 26.5). Uses the long-standing +/// `addCompletedHandler` callback bridged through +/// `withCheckedContinuation`, which works on every Metal SDK from +/// macOS 10.11 onward. +extension MTLCommandBuffer { + public func awaitCompletion() async { + await withCheckedContinuation { (cont: CheckedContinuation) in + self.addCompletedHandler { _ in cont.resume() } + } + } +} // MARK: - ModelLoadLock diff --git a/Tests/MetalTileSwiftTests/PSOCacheTests.swift b/Tests/MetalTileSwiftTests/PSOCacheTests.swift index d4002b6a..187b66f9 100644 --- a/Tests/MetalTileSwiftTests/PSOCacheTests.swift +++ b/Tests/MetalTileSwiftTests/PSOCacheTests.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -24,7 +24,7 @@ // lookup in `PSOCache.lookup`). import Foundation -import Metal +@preconcurrency import Metal import Testing @testable import MetalTileSwift diff --git a/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift b/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift new file mode 100644 index 00000000..5b08a069 --- /dev/null +++ b/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift @@ -0,0 +1,286 @@ +// Copyright 2026 Tom Turney (@TheTom) +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// Qwen3.5-35B-A3B (MoE + GDN hybrid) forwardMany / decode perf bench. +// +// Same shape as `Qwen36TextIntegrationTests.forwardManyBench*` but +// pointed at the locally-available Qwen3.5-35B-A3B-4bit checkpoint. +// Identical engine path (`m.qwen35`), so this exercises every Bagel +// wiring landed on `tom/bagel-clean`: GPU MoE router, batched QKV +// fast paths, GDN chunked recurrence, rmsNormQgemvInt4Fast finalNorm, +// 2-pass Flash-Decoding, MTLResidencySet weight pinning, etc. + +import Foundation +import TestHelpers +import Testing + +@testable import FFAI + +private let qwen35MoELocalPath = "/Users/tom/models/Qwen3.5-35B-A3B-4bit" + +@Suite("Qwen3.5-35B-A3B local-checkpoint bench", .serialized) +struct Qwen35MoEBenchIntegrationTests { + + @Test("Qwen3.5-35B-A3B decode T=1 tps — 5 runs, median over 32 steps") + func decodeBenchT1() async throws { + guard FileManager.default.fileExists(atPath: qwen35MoELocalPath) else { + print("decodeBenchT1 skipped: \(qwen35MoELocalPath) not found") + return + } + let m = try await loadModel() + guard let qwen = m.qwen35 else { + Issue.record("expected Qwen35Model engine") + return + } + + let prompt = "The history of the printing press began when" + let promptTokens = m.tokenizer.encode(text: prompt) + let promptLen = promptTokens.count + + // Warm. + for _ in 0 ..< 2 { + let warmCaches = qwen.makeLayerCaches() + for (i, tok) in promptTokens.enumerated() { + _ = qwen.forward(tokenId: tok, position: i, caches: warmCaches) + } + for j in 0 ..< 4 { + _ = qwen.forward( + tokenId: 0, position: promptLen + j, caches: warmCaches) + } + } + + // 5 timed runs of decode-only, 32 steps each (median). + let nSteps = 32 + var runs: [Double] = [] + for _ in 0 ..< 5 { + let caches = qwen.makeLayerCaches() + for (i, tok) in promptTokens.enumerated() { + _ = qwen.forward(tokenId: tok, position: i, caches: caches) + } + let t0 = Date() + for j in 0 ..< nSteps { + _ = qwen.forward( + tokenId: 0, position: promptLen + j, caches: caches) + } + runs.append(Date().timeIntervalSince(t0)) + } + runs.sort() + let median = runs[runs.count / 2] + let tps = Double(nSteps) / median + print( + "Qwen3.5-35B-A3B decode T=1: runs=\(runs.map { String(format: "%.3f", $0) })s " + + "median=\(String(format: "%.3f", median))s → \(String(format: "%.2f", tps)) tps" + ) + } + + @Test("Qwen3.5-35B-A3B forwardMany T=8 smoke") + func forwardManyT8Smoke() async throws { + guard FileManager.default.fileExists(atPath: qwen35MoELocalPath) else { + print("forwardManyT8Smoke skipped: \(qwen35MoELocalPath) not found") + return + } + let m = try await loadModel() + let qwen = try #require(m.qwen35, "expected Qwen35Model engine") + let seed = "The quick brown fox" + var encoded = m.tokenizer.encode(text: seed) + while encoded.count < 8 { encoded.append(0) } + encoded = Array(encoded.prefix(8)) + let caches = qwen.makeLayerCaches() + let cmd = Device.shared.makeCommandBuffer() + let t0 = Date() + _ = qwen.forwardMany( + tokenIds: encoded, startPosition: 0, + caches: caches, on: cmd, device: Device.shared) + cmd.commit() + await cmd.awaitCompletion() + let dt = Date().timeIntervalSince(t0) + print("Qwen3.5-35B-A3B forwardMany T=8: \(String(format: "%.3f", dt))s") + } + + @Test("Qwen3.5-35B-A3B per-token forward T=128 (no batched)") + func perTokenT128() async throws { + guard FileManager.default.fileExists(atPath: qwen35MoELocalPath) else { + print("perTokenT128 skipped: \(qwen35MoELocalPath) not found") + return + } + let m = try await loadModel() + let qwen = try #require(m.qwen35, "expected Qwen35Model engine") + let seed = "The quick brown fox jumps over the lazy dog. " + let seedEncoded = m.tokenizer.encode(text: seed) + var encoded = seedEncoded + while encoded.count < 128 { encoded.append(contentsOf: seedEncoded) } + encoded = Array(encoded.prefix(128)) + let caches = qwen.makeLayerCaches() + let t0 = Date() + for (i, tok) in encoded.enumerated() { + _ = qwen.forward(tokenId: tok, position: i, caches: caches) + } + let dt = Date().timeIntervalSince(t0) + print( + "Qwen3.5-35B-A3B per-token T=128: \(String(format: "%.3f", dt))s → " + + "\(String(format: "%.1f", 128.0 / dt)) tps") + } + + @Test("Qwen3.5-35B-A3B forwardManyBench T=128") + func forwardManyT128() async throws { + try await runForwardManyBench(targetT: 128) + } + + @Test("Qwen3.5-35B-A3B forwardManyBench T=512") + func forwardManyT512() async throws { + try await runForwardManyBench(targetT: 512) + } + + @Test("Qwen3.5-35B-A3B forwardManyBench T=2048 (long-context)") + func forwardManyT2K() async throws { + try await runForwardManyBench(targetT: 2048) + } + + @Test("Qwen3.5-35B-A3B prefill T=32K + decode-after-prefill — long-context bench") + func longContext32K() async throws { + guard FileManager.default.fileExists(atPath: qwen35MoELocalPath) else { + print("longContext32K skipped: \(qwen35MoELocalPath) not found") + return + } + let m = try await loadModel() + let qwen = try #require(m.qwen35, "expected Qwen35Model engine") + let seed = + "The history of the printing press began when European craftsmen of the 15th century combined movable metal type with oil based ink screw presses paper to mass produce printed books pamphlets and broadsheets revolutionising communication" + let seedEncoded = m.tokenizer.encode(text: seed) + let targetT = 32_768 + var encoded = seedEncoded + while encoded.count < targetT { + encoded.append(contentsOf: seedEncoded) + } + encoded = Array(encoded.prefix(targetT)) + let T = encoded.count + + // Single prefill + decode run (32 decode steps). No + // 5-run median — the prefill alone is multi-second so the + // variance from a single shot is fine for a sanity bench. + // Per-token baseline at T=32K is intentionally NOT measured + // (would take ~5 minutes per run). + print("Qwen3.5-35B-A3B longContext32K T=\(T) (single run)") + + let nDecode = 32 + let caches = qwen.makeLayerCaches() + let prefillT0 = Date() + let cmd = Device.shared.makeCommandBuffer() + _ = qwen.forwardMany( + tokenIds: encoded, startPosition: 0, + caches: caches, on: cmd, device: Device.shared) + cmd.commit() + await cmd.awaitCompletion() + let prefillS = Date().timeIntervalSince(prefillT0) + let prefillTps = Double(T) / prefillS + print( + "Qwen3.5-35B-A3B prefill T=\(T): " + + "\(String(format: "%.3f", prefillS))s → " + + "\(String(format: "%.1f", prefillTps)) tps batched") + + let decodeT0 = Date() + for j in 0 ..< nDecode { + _ = qwen.forward( + tokenId: 0, position: T + j, caches: caches) + } + let decodeS = Date().timeIntervalSince(decodeT0) + let decodeTps = Double(nDecode) / decodeS + print( + "Qwen3.5-35B-A3B decode T=1 after T=\(T) prefill: " + + "\(String(format: "%.3f", decodeS))s over \(nDecode) steps → " + + "\(String(format: "%.2f", decodeTps)) tps") + } + + private func loadModel() async throws -> Model { + var optsBuilder = LoadOptions() + optsBuilder.prewarm = false + let opts = optsBuilder + return try await ModelLoadLock.shared.loadSerially { + try await Model.load(qwen35MoELocalPath, options: opts) + } + } + + private func runForwardManyBench(targetT: Int) async throws { + guard FileManager.default.fileExists(atPath: qwen35MoELocalPath) else { + print("forwardManyBench skipped: \(qwen35MoELocalPath) not found") + return + } + let m = try await loadModel() + let qwen = try #require(m.qwen35, "expected Qwen35Model engine") + + let seed = + "The history of the printing press began when European craftsmen of the 15th century combined movable metal type with oil based ink screw presses paper to mass produce printed books pamphlets and broadsheets revolutionising communication" + let seedEncoded = m.tokenizer.encode(text: seed) + var encoded = seedEncoded + while encoded.count < targetT { + encoded.append(contentsOf: seedEncoded) + } + encoded = Array(encoded.prefix(targetT)) + let T = encoded.count + print("Qwen3.5-35B-A3B forwardManyBench T=\(T)") + + // Warm. + for _ in 0 ..< 2 { + let warmCachesP = qwen.makeLayerCaches() + for (i, tok) in encoded.prefix(2).enumerated() { + _ = qwen.forward(tokenId: tok, position: i, caches: warmCachesP) + } + let warmCachesB = qwen.makeLayerCaches() + let warmCmd = Device.shared.makeCommandBuffer() + _ = qwen.forwardMany( + tokenIds: encoded, startPosition: 0, + caches: warmCachesB, on: warmCmd, device: Device.shared) + warmCmd.commit() + await warmCmd.awaitCompletion() + } + + // Per-token loop baseline (5 runs, median). + var perTokenSecs: [Double] = [] + for _ in 0 ..< 5 { + let caches = qwen.makeLayerCaches() + let t0 = Date() + for (i, tok) in encoded.enumerated() { + _ = qwen.forward(tokenId: tok, position: i, caches: caches) + } + perTokenSecs.append(Date().timeIntervalSince(t0)) + } + perTokenSecs.sort() + let perTokenMedian = perTokenSecs[perTokenSecs.count / 2] + + // Batched forwardMany (5 runs, median). + var batchedSecs: [Double] = [] + for _ in 0 ..< 5 { + let caches = qwen.makeLayerCaches() + let bCmd = Device.shared.makeCommandBuffer() + let t0 = Date() + _ = qwen.forwardMany( + tokenIds: encoded, startPosition: 0, + caches: caches, on: bCmd, device: Device.shared) + bCmd.commit() + await bCmd.awaitCompletion() + batchedSecs.append(Date().timeIntervalSince(t0)) + } + batchedSecs.sort() + let batchedMedian = batchedSecs[batchedSecs.count / 2] + + let speedup = perTokenMedian / batchedMedian + let batchedTps = Double(T) / batchedMedian + print( + "Qwen3.5-35B-A3B RESULT T=\(T): " + + "per_token=\(String(format: "%.0f", perTokenMedian * 1000))ms " + + "batched=\(String(format: "%.0f", batchedMedian * 1000))ms " + + "speedup=\(String(format: "%.2fx", speedup)) " + + "batched_tps=\(String(format: "%.1f", batchedTps))") + } +} diff --git a/Tests/ModelIntegrationTests/Text/Qwen36TextIntegrationTests.swift b/Tests/ModelIntegrationTests/Text/Qwen36TextIntegrationTests.swift index e97f28ec..fbee50ad 100644 --- a/Tests/ModelIntegrationTests/Text/Qwen36TextIntegrationTests.swift +++ b/Tests/ModelIntegrationTests/Text/Qwen36TextIntegrationTests.swift @@ -1,4 +1,4 @@ -// Copyright 2026 Eric Kryski (@ekryski) +// Copyright 2026 Eric Kryski (@ekryski) and Tom Turney (@TheTom) // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -179,6 +179,11 @@ struct Qwen36TextIntegrationTests { try await runForwardManyBench(targetT: 512) } + @Test("Qwen3.6-35B-A3B forwardMany bench — T=2048 long-context scaling") + func forwardManyBench2K() async throws { + try await runForwardManyBench(targetT: 2048) + } + private func runForwardManyBench(targetT: Int) async throws { let path = qwen36LocalPath var optsBuilder = LoadOptions() @@ -215,7 +220,7 @@ struct Qwen36TextIntegrationTests { tokenIds: encoded, startPosition: 0, caches: warmCachesB, on: warmCmd, device: Device.shared) warmCmd.commit() - await warmCmd.completed() + await warmCmd.awaitCompletion() _ = warmIter // silence } @@ -245,7 +250,7 @@ struct Qwen36TextIntegrationTests { tokenIds: encoded, startPosition: 0, caches: caches, on: bCmd, device: Device.shared) bCmd.commit() - await bCmd.completed() + await bCmd.awaitCompletion() batchedSecs.append(Date().timeIntervalSince(t0)) } batchedSecs.sort() @@ -292,7 +297,7 @@ struct Qwen36TextIntegrationTests { tokenIds: encoded, startPosition: 0, caches: manyCaches, on: manyCmd, device: Device.shared) manyCmd.commit() - await manyCmd.completed() + await manyCmd.awaitCompletion() let manyLogits = manyLogitsTensor.toFloatArray() let manyArgmax = manyLogits.enumerated().max(by: { $0.element < $1.element })!.offset diff --git a/Tests/ModelIntegrationTests/Text/SpecDecodeBenchTests.swift b/Tests/ModelIntegrationTests/Text/SpecDecodeBenchTests.swift new file mode 100644 index 00000000..e524215a --- /dev/null +++ b/Tests/ModelIntegrationTests/Text/SpecDecodeBenchTests.swift @@ -0,0 +1,174 @@ +// Copyright 2026 Tom Turney (@TheTom) +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// End-to-end spec-decode bench. Compares: +// * Baseline greedy decode (32 steps via Qwen35Model.forward) +// * SpecDecode.generateGreedy with NGramDrafter (γ ∈ {1, 2, 4}) +// +// Verifies generated sequences are IDENTICAL (greedy is deterministic +// so spec decode must produce the same tokens when correctly +// implemented) and reports decode-only tps for each path so the cost +// of the verify forward is visible. + +import Foundation +import Metal +import TestHelpers +import Testing + +@testable import FFAI + +private let qwen36SpecPath = "/Users/tom/models/Qwen3.6-35B-A3B-4bit" + +@Suite("SpecDecode end-to-end vs baseline greedy") +struct SpecDecodeBenchTests { + + @Test("SpecDecode + NGramDrafter produces same tokens as baseline greedy, with measurable tps") + func specDecodeMatchesBaselineGreedyTPSReport() async throws { + guard FileManager.default.fileExists(atPath: qwen36SpecPath) else { + print("SpecDecodeBench skipped: \(qwen36SpecPath) not found") + return + } + var optsBuilder = LoadOptions() + optsBuilder.prewarm = false + let opts = optsBuilder + let m: Model = try await ModelLoadLock.shared.loadSerially { + try await Model.load(qwen36SpecPath, options: opts) + } + guard let qwen = m.qwen35 else { + Issue.record("expected Qwen35Model engine") + return + } + + // Prompt designed to be repetitive enough that the n-gram + // drafter can hit. Code patterns are the sweet spot — pure + // prose generation has low n-gram acceptance. + let prompt = """ + def fibonacci(n): + if n <= 1: + return n + return fibonacci(n - 1) + fibonacci(n - 2) + + def fibonacci_iterative(n): + if n <= 1: + return n + a, b = 0, 1 + for i in range(2, n + 1): + a, b = b, a + b + return b + + def + """ + let promptTokens = m.tokenizer.encode(text: prompt) + let promptLen = promptTokens.count + precondition(promptLen >= 4, "need at least 4 prompt tokens") + print("SpecDecodeBench prompt=\(promptLen) tokens") + let maxNewTokens = 32 + + // ── Baseline: greedy decode via forward() loop ──────────────── + var baselineTokens: [Int] = [] + let baseCaches = qwen.makeLayerCaches() + // Prefill (untimed; same cost for both paths). + var lastLogits: Tensor! + for (i, tok) in promptTokens.enumerated() { + lastLogits = qwen.forward( + tokenId: tok, position: i, caches: baseCaches) + } + var pos = promptLen + var next = argmaxHost(lastLogits) + // Decode loop — TIMED region (decode-only, no prefill). + let baseDecodeT0 = Date() + for _ in 0 ..< maxNewTokens { + baselineTokens.append(next) + let logits = qwen.forward( + tokenId: next, position: pos, caches: baseCaches) + next = argmaxHost(logits) + pos += 1 + } + let baseDecodeS = Date().timeIntervalSince(baseDecodeT0) + let baseTps = Double(maxNewTokens) / baseDecodeS + + // ── SpecDecode + NGramDrafter ───────────────────────────────── + for gamma in [1, 2, 4] { + let specCaches = qwen.makeLayerCaches() + var history = promptTokens + // Prefill (untimed). + var prefillLastLogits: Tensor! + for (i, tok) in promptTokens.enumerated() { + prefillLastLogits = qwen.forward( + tokenId: tok, position: i, caches: specCaches) + } + let prefillFirstSample = argmaxHost(prefillLastLogits) + + // Spec decode loop — TIMED region. + let drafter = NGramDrafter(maxNMatch: 3, minNMatch: 2) + let specT0 = Date() + let stats = SpecDecode.generateGreedy( + model: qwen, + drafter: drafter, + gamma: gamma, + lastToken: prefillFirstSample, + position: promptLen, + caches: specCaches, + history: &history, + maxNewTokens: maxNewTokens) + let specS = Date().timeIntervalSince(specT0) + + // Generated tokens are history[promptLen..<...]. A full- + // accept iter can overshoot maxNewTokens by up to γ, so + // clamp from the prompt boundary, not from the end. + let genEnd = Swift.min(history.count, promptLen + maxNewTokens) + let specGenerated = Array(history[promptLen ..< genEnd]) + + let matchPrefix = zip(baselineTokens, specGenerated) + .prefix(while: { $0 == $1 }).count + let specTps = Double(maxNewTokens) / specS + let speedup = specTps / baseTps + print( + "SpecDecode γ=\(gamma): decode_only=\(String(format: "%.3f", specS))s, " + + "tps=\(String(format: "%.2f", specTps)), " + + "vs_baseline=\(String(format: "%.2fx", speedup)), " + + "accepted=\(stats.candidatesAccepted)/\(stats.candidatesProposed) " + + "(\(String(format: "%.1f", stats.acceptanceRate * 100))%), " + + "fallback_steps=\(stats.fallbackSingleSteps), " + + "matches_baseline_prefix=\(matchPrefix)/\(maxNewTokens)") + + // Greedy + greedy drafter + correct snapshot/restore MUST + // match baseline exactly. Print first-8-token diff on + // mismatch so the driver can be debugged without + // hard-failing while it's settling. + if matchPrefix < maxNewTokens { + print(" baseline: \(baselineTokens.prefix(8))") + print(" spec : \(specGenerated.prefix(8))") + } + } + + print( + "Baseline greedy decode-only: \(String(format: "%.3f", baseDecodeS))s, tps=\(String(format: "%.2f", baseTps)) over \(maxNewTokens) steps" + ) + } +} + +@inline(__always) +private func argmaxHost(_ logits: Tensor) -> Int { + let host = logits.toFloatArray() + var bestIdx = 0 + var bestVal = host[0] + for i in 1 ..< host.count { + if host[i] > bestVal { + bestVal = host[i] + bestIdx = i + } + } + return bestIdx +} diff --git a/Tests/ModelIntegrationTests/Text/SpecDecodeVerifyCostBench.swift b/Tests/ModelIntegrationTests/Text/SpecDecodeVerifyCostBench.swift new file mode 100644 index 00000000..3f06325d --- /dev/null +++ b/Tests/ModelIntegrationTests/Text/SpecDecodeVerifyCostBench.swift @@ -0,0 +1,189 @@ +// Copyright 2026 Tom Turney (@TheTom) +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SpecDecodeVerifyCostBench — answers: at γ+1 ≤ 4, does the optimised +// decode T=1 path (which benefits from scalarFMA + bm8 qmm + GDN +// fused prep) beat the batched `forwardManyAllLogits` path? +// +// Methodology: time N iterations of each shape on identical caches; +// `snapshotAll` / `restoreAll` between iters so we measure pure +// forward work without context drift. Reports median ms / iter and +// the batched-vs-single ratio for both T=3 (γ=2 verify) and T=2 +// (γ=1 verify). + +import Foundation +import Metal +import TestHelpers +import Testing + +@testable import FFAI + +private let qwen36VerifyPath = "/Users/tom/models/Qwen3.6-35B-A3B-4bit" + +@Suite("SpecDecode verify-cost bench") +struct SpecDecodeVerifyCostBench { + + @Test("Compare forwardManyAllLogits(T=3) vs 3× forward() at decode shape") + func compareBatchedVsSingleStep() async throws { + guard FileManager.default.fileExists(atPath: qwen36VerifyPath) else { + print("SpecDecodeVerifyCostBench skipped: \(qwen36VerifyPath) not found") + return + } + var optsBuilder = LoadOptions() + optsBuilder.prewarm = false + let opts = optsBuilder + let m: Model = try await ModelLoadLock.shared.loadSerially { + try await Model.load(qwen36VerifyPath, options: opts) + } + guard let qwen = m.qwen35 else { + Issue.record("expected Qwen35Model engine") + return + } + + let prompt = """ + def fibonacci(n): + if n <= 1: + return n + return fibonacci(n - 1) + fibonacci(n - 2) + + def + """ + let promptTokens = m.tokenizer.encode(text: prompt) + let promptLen = promptTokens.count + let device = Device.shared + + let caches = qwen.makeLayerCaches() + // Prefill — untimed. + for (i, tok) in promptTokens.enumerated() { + _ = qwen.forward(tokenId: tok, position: i, caches: caches) + } + + let inputIds = [ + promptTokens[promptLen - 4], + promptTokens[promptLen - 3], + promptTokens[promptLen - 2], + ] + + // Snapshot caches so iteration N starts in identical state. + let snap0 = caches.snapshotAll(device: device) + + // Warm the single-step path. + for tok in inputIds { + _ = qwen.forward(tokenId: tok, position: promptLen, caches: caches) + } + caches.restoreAll(from: snap0, device: device) + + // Warm the batched path. + let warmCmd = device.makeCommandBuffer() + _ = qwen.forwardManyAllLogits( + tokenIds: inputIds, startPosition: promptLen, + caches: caches, on: warmCmd, device: device) + warmCmd.commit() + await warmCmd.awaitCompletion() + caches.restoreAll(from: snap0, device: device) + + // 16 timed iters of each path, restoring caches between each. + let nIters = 16 + + var singleTimes: [Double] = [] + for _ in 0 ..< nIters { + let t0 = Date() + for (i, tok) in inputIds.enumerated() { + _ = qwen.forward( + tokenId: tok, position: promptLen + i, caches: caches) + } + singleTimes.append(Date().timeIntervalSince(t0)) + caches.restoreAll(from: snap0, device: device) + } + + var batchedTimes: [Double] = [] + for _ in 0 ..< nIters { + let t0 = Date() + let cmd = device.makeCommandBuffer() + _ = qwen.forwardManyAllLogits( + tokenIds: inputIds, startPosition: promptLen, + caches: caches, on: cmd, device: device) + cmd.commit() + await cmd.awaitCompletion() + batchedTimes.append(Date().timeIntervalSince(t0)) + caches.restoreAll(from: snap0, device: device) + } + + let singleMedian = singleTimes.sorted()[nIters / 2] + let batchedMedian = batchedTimes.sorted()[nIters / 2] + let singleMs = singleMedian * 1000 + let batchedMs = batchedMedian * 1000 + let ratio = batchedMs / singleMs + + print("VerifyCostBench T=3 (median of \(nIters) iters):") + print(" 3× forward() loop: \(String(format: "%.2f", singleMs)) ms") + print(" 1× forwardManyAllLogits: \(String(format: "%.2f", batchedMs)) ms") + print(" batched / single ratio: \(String(format: "%.2fx", ratio))") + if ratio > 1.0 { + print( + " → single-step LOOP is faster by \(String(format: "%.1f", (ratio - 1.0) * 100))%" + ) + } else { + print( + " → batched is faster by \(String(format: "%.1f", (1.0 / ratio - 1.0) * 100))%" + ) + } + + // T=2 case (γ=1 verify shape). + let inputIds2 = [ + promptTokens[promptLen - 2], promptTokens[promptLen - 1], + ] + let snap2 = caches.snapshotAll(device: device) + for tok in inputIds2 { + _ = qwen.forward(tokenId: tok, position: promptLen, caches: caches) + } + caches.restoreAll(from: snap2, device: device) + let warmCmd2 = device.makeCommandBuffer() + _ = qwen.forwardManyAllLogits( + tokenIds: inputIds2, startPosition: promptLen, + caches: caches, on: warmCmd2, device: device) + warmCmd2.commit() + await warmCmd2.awaitCompletion() + caches.restoreAll(from: snap2, device: device) + + var singleT2: [Double] = [] + for _ in 0 ..< nIters { + let t0 = Date() + for (i, tok) in inputIds2.enumerated() { + _ = qwen.forward( + tokenId: tok, position: promptLen + i, caches: caches) + } + singleT2.append(Date().timeIntervalSince(t0)) + caches.restoreAll(from: snap2, device: device) + } + var batchedT2: [Double] = [] + for _ in 0 ..< nIters { + let t0 = Date() + let cmd = device.makeCommandBuffer() + _ = qwen.forwardManyAllLogits( + tokenIds: inputIds2, startPosition: promptLen, + caches: caches, on: cmd, device: device) + cmd.commit() + await cmd.awaitCompletion() + batchedT2.append(Date().timeIntervalSince(t0)) + caches.restoreAll(from: snap2, device: device) + } + let s2ms = singleT2.sorted()[nIters / 2] * 1000 + let b2ms = batchedT2.sorted()[nIters / 2] * 1000 + print("VerifyCostBench T=2 (γ=1 shape, median \(nIters) iters):") + print(" 2× forward() loop: \(String(format: "%.2f", s2ms)) ms") + print(" 1× forwardManyAllLogits: \(String(format: "%.2f", b2ms)) ms") + print(" ratio: \(String(format: "%.2fx", b2ms / s2ms))") + } +}