From 2eab5a1cb9b1501387490bb52790beeef8390233 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 07:09:25 -0500
Subject: [PATCH 01/54] perf(device): MTLResidencySet pin helper on Device

`Device.markWeightsResident(_:)` attaches the given MTLBuffers to a
persistent residency set on the command queue. Without it, the Metal
driver re-validates per-allocation residency on every command-buffer
encode; at large MoE models with thousands of dispatches per prefill,
that overhead dominates wall time.

The set is shared across all weight buffers and lazily created on
first call. Repeated calls add to it. macOS 15 / iOS 18 minimum;
older OSes no-op silently. `FFAI_NO_RESIDENCY_SET=1` disables for
A/B benching.

Per-model integration (call `markWeightsResident` after `Model.load`)
lands in a follow-up commit so this stays a pure infrastructure
add. KV / GDN / conv state cache pinning also lands separately so
the residency-set adoption can be diffed per surface.
---
 Sources/FFAI/Device.swift | 46 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/Sources/FFAI/Device.swift b/Sources/FFAI/Device.swift
index c63d9899..d6838c10 100644
--- a/Sources/FFAI/Device.swift
+++ b/Sources/FFAI/Device.swift
@@ -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()
+    }
 }

From a47f0779a7b4477d8cb6635ec2a4c197b746f62e Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 07:16:11 -0500
Subject: [PATCH 02/54] perf(loader): pin model weights in MTLResidencySet at
 load time

Walks every entry of the loaded `SafeTensorsBundle` and registers its
backing MTLBuffer with `Device.markWeightsResident(_:)` immediately
after the bundle resolves, before `ModelRegistry.dispatchAndLoad` kicks
the per-family loader. The driver's per-command-buffer residency
re-validation is the cost we avoid.

Opt-out via `FFAI_NO_RESIDENCY_SET=1` (handled inside the helper).
Older OSes silently no-op via the `@available(macOS 15, iOS 18, *)`
gate on the helper.
---
 Sources/FFAI/Loader/Model.swift | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/Sources/FFAI/Loader/Model.swift b/Sources/FFAI/Loader/Model.swift
index f86d6c94..fb0eba39 100644
--- a/Sources/FFAI/Loader/Model.swift
+++ b/Sources/FFAI/Loader/Model.swift
@@ -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
                 )

From aa833d76d869217057e3ecdf2fd19f8acaf5589a Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 07:20:47 -0500
Subject: [PATCH 03/54] KVCache: pin K/V/GDN/conv buffers in residency set on
 init
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Each cache class owns one or two MTLBuffers that persist for the
lifetime of a generation and get touched by every decode-step
dispatch. Without explicit residency pinning, Apple's Metal driver
re-validates per-allocation residency on every command-buffer
encode — at decode rates of thousands of dispatches per token, the
per-dispatch overhead dominates wall time.

Reuses the device-wide MTLResidencySet introduced for weight
buffers (see `Device.markWeightsResident`). The set is single-flighted
under a lock and shared across all cache classes; the underlying
Metal API gracefully no-ops on macOS < 15 and is gated behind
`FFAI_NO_RESIDENCY_SET=1` for A/B comparison.

- `KVCache.init` pins kBuffer + vBuffer
- `GDNStateCache.init` pins current + next recurrent-state buffers
- `ConvStateCache.init` pins the rolling-window state

Migrated from PR #9 (Everything Bagel 1) — measured win was part of
the 16.29 → 55.94 tps decode T=1 jump on Qwen3.6-A3B / M5 Max.
---
 Sources/FFAI/KVCache/ConvStateCache.swift | 4 ++++
 Sources/FFAI/KVCache/GDNStateCache.swift  | 5 +++++
 Sources/FFAI/KVCache/KVCache.swift        | 4 ++++
 3 files changed, 13 insertions(+)

diff --git a/Sources/FFAI/KVCache/ConvStateCache.swift b/Sources/FFAI/KVCache/ConvStateCache.swift
index b9cbfccc..3a78a0f6 100644
--- a/Sources/FFAI/KVCache/ConvStateCache.swift
+++ b/Sources/FFAI/KVCache/ConvStateCache.swift
@@ -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).
diff --git a/Sources/FFAI/KVCache/GDNStateCache.swift b/Sources/FFAI/KVCache/GDNStateCache.swift
index f0162591..b96dca8b 100644
--- a/Sources/FFAI/KVCache/GDNStateCache.swift
+++ b/Sources/FFAI/KVCache/GDNStateCache.swift
@@ -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`
diff --git a/Sources/FFAI/KVCache/KVCache.swift b/Sources/FFAI/KVCache/KVCache.swift
index 43e46a25..97a81a3e 100644
--- a/Sources/FFAI/KVCache/KVCache.swift
+++ b/Sources/FFAI/KVCache/KVCache.swift
@@ -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

From 50b45e44cc3ae182eb33f3eef03520eee56a60aa Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 07:26:47 -0500
Subject: [PATCH 04/54] Ops: add scalarFMA family + sigmoidScalarFMAResidual
 for MoE accumulator
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The Qwen3-MoE decode path needs a scalar-broadcast FMA primitive to
accumulate the top-K expert outputs into a single hidden-state
vector. Without it the path was forced into a
`Tensor.filled([hidden], weight)` host alloc per expert plus a
`mul + add` pair per expert — a host roundtrip and 2 dispatches the
GPU could absorb into one.

Adds four wrappers (all in `Ops/OpsFused.swift`) over kernels that
metaltile already emits:

- `scalarFMA(scalar, value, base, into:)` — single-dispatch
  `out[i] = base[i] + scalar[0] * value[i]`.
- `scalarFMAMany(scalars, values, acc:)` — N sequential dispatches
  reusing one compute encoder; Metal guarantees in-order encoder
  execution so the accumulation is safe.
- `scalarFMAChain8(scalars, values, out:)` — exact 8-way fused
  dispatch (one kernel reads 16 inputs at once), collapses the
  topK=8 expert accumulator into a single dispatch.
- `sigmoidScalarFMAResidual(gate, value, base, residual, into:)` —
  sigmoid-gated FMA with the post-MoE residual add folded in;
  Eric's `sigmoidScalarFMA` covers the non-residual variant.

Tests in `Tests/FFAITests/Ops/OpsSpecialPathTests.swift` cover each
wrapper with a closed-form expected value (4 new tests).

Migrated from PR #9 / #10 (Everything Bagel 1 + 2) — measured as
load-bearing for the decode T=1 16.29 → 96.3 tps win on
Qwen3.6-A3B / M5 Max.
---
 Sources/FFAI/Ops/OpsFused.swift               | 264 ++++++++++++++++++
 Tests/FFAITests/Ops/OpsSpecialPathTests.swift | 117 ++++++++
 2 files changed, 381 insertions(+)

diff --git a/Sources/FFAI/Ops/OpsFused.swift b/Sources/FFAI/Ops/OpsFused.swift
index c6dc7f80..2aa5f3b4 100644
--- a/Sources/FFAI/Ops/OpsFused.swift
+++ b/Sources/FFAI/Ops/OpsFused.swift
@@ -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,261 @@ 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)")
+        }
+    }
+
+    /// `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/Tests/FFAITests/Ops/OpsSpecialPathTests.swift b/Tests/FFAITests/Ops/OpsSpecialPathTests.swift
index 3b2dc65f..2ef6836e 100644
--- a/Tests/FFAITests/Ops/OpsSpecialPathTests.swift
+++ b/Tests/FFAITests/Ops/OpsSpecialPathTests.swift
@@ -128,6 +128,123 @@ 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: - KV cache append
 
     @Test("kvCacheUpdate f32 — writes one row into [nKV, maxSeq, headDim]")

From 344a56552549ffc463da96d8ef76db55f6b06b6e Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 07:29:28 -0500
Subject: [PATCH 05/54] =?UTF-8?q?Ops:=20add=20dequantGemvInt4Two/Three/Fou?=
 =?UTF-8?q?r=20=E2=80=94=20batched=20int4=20dequantGemv?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Three batched wrappers around the existing `dequant_gemv_int4_*`
metaltile kernels that drive N back-to-back dispatches inside a
single compute encoder, sharing the `input` binding and the
`in_dim` / `group_size` constants. Saves N-1 encoder begin/end
pairs per call.

- `dequantGemvInt4Two`  — per-expert MoE gate + up projections.
- `dequantGemvInt4Three` — Qwen3-style attention q/k/v projections
  in the unfused decode path.
- `dequantGemvInt4Four`  — Qwen3.5 GDN mixer's qkv/z/b/a
  input-projection quartet, all reading the same xNorm.

CPU-correctness tests in `Tests/FFAITests/Ops/QuantizedOpsTests.swift`
reuse a `makeInt4Projection` helper to build packed weights +
scales + biases + expected GEMV output, then verify each batched
wrapper matches CPU truth for every output tensor.

Migrated from PR #9 / #10.
---
 Sources/FFAI/Ops/Ops.swift                  | 156 ++++++++++++++++++
 Tests/FFAITests/Ops/QuantizedOpsTests.swift | 167 ++++++++++++++++++++
 2 files changed, 323 insertions(+)

diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift
index ec4695be..60f55393 100644
--- a/Sources/FFAI/Ops/Ops.swift
+++ b/Sources/FFAI/Ops/Ops.swift
@@ -1319,6 +1319,162 @@ public enum Ops {
             on: cmd, into: out)
     }
 
+    /// 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(
+            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()
+    }
+
     /// 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.
diff --git a/Tests/FFAITests/Ops/QuantizedOpsTests.swift b/Tests/FFAITests/Ops/QuantizedOpsTests.swift
index a2b2e2bd..4b3739d5 100644
--- a/Tests/FFAITests/Ops/QuantizedOpsTests.swift
+++ b/Tests/FFAITests/Ops/QuantizedOpsTests.swift
@@ -702,4 +702,171 @@ 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)
+            }
+        }
+    }
 }

From 49871a93cec9a5706c4b4b716ca3065e49d12990 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 07:32:05 -0500
Subject: [PATCH 06/54] =?UTF-8?q?Ops:=20add=20ropePartialTwo=20=E2=80=94?=
 =?UTF-8?q?=20pair=20Q+K=20RoPE=20in=20one=20compute=20encoder?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Both Q and K go through the same partial RoPE rotation with identical
`(headDim, rotaryDim, position, thetaBase, scaling)` parameters
every decode step on Qwen3-class models. `ropePartialTwo` sets the
RoPE constants once and dispatches both rotations through a single
shared compute encoder, saving one encoder begin/end pair per
attention layer per token.

The kernel is unchanged (`ffai_rope_llama_*`); the wrapper is a
buffer-binding rotation around the existing in-place rotation. A
test verifies bit-identical output against two sequential
`ropePartial` calls on the same inputs.

Migrated from PR #10.
---
 Sources/FFAI/Ops/Ops.swift         | 63 ++++++++++++++++++++++++++++++
 Tests/FFAITests/Ops/OpsTests.swift | 39 ++++++++++++++++++
 2 files changed, 102 insertions(+)

diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift
index 60f55393..f8b7cc90 100644
--- a/Sources/FFAI/Ops/Ops.swift
+++ b/Sources/FFAI/Ops/Ops.swift
@@ -840,6 +840,69 @@ 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()
+    }
+
     /// 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
diff --git a/Tests/FFAITests/Ops/OpsTests.swift b/Tests/FFAITests/Ops/OpsTests.swift
index a45466e3..bc6a9d8e 100644
--- a/Tests/FFAITests/Ops/OpsTests.swift
+++ b/Tests/FFAITests/Ops/OpsTests.swift
@@ -110,6 +110,45 @@ 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("ropePartial f32 — rotaryDim == headDim matches full rope")
     func ropePartialFullEqualsRope() {
         autoreleasepool {

From ec295535ba0a8fc25a5af948df50f8539352e920 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 07:34:12 -0500
Subject: [PATCH 07/54] =?UTF-8?q?Ops:=20add=20gatherTwo=20=E2=80=94=20two?=
 =?UTF-8?q?=20embedding=20lookups=20on=20one=20encoder?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Pairs two `ffai_gather_*` dispatches against the same `table`
binding inside one compute encoder, saving an encoder begin/end
pair when the caller needs to fetch two parallel id streams from
the same embedding (e.g. the Qwen3 attention head-half split path).

Migrated from PR #10.
---
 Sources/FFAI/Ops/Ops.swift         | 48 ++++++++++++++++++++++++++++++
 Tests/FFAITests/Ops/OpsTests.swift | 23 ++++++++++++++
 2 files changed, 71 insertions(+)

diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift
index f8b7cc90..4564df7b 100644
--- a/Sources/FFAI/Ops/Ops.swift
+++ b/Sources/FFAI/Ops/Ops.swift
@@ -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.
diff --git a/Tests/FFAITests/Ops/OpsTests.swift b/Tests/FFAITests/Ops/OpsTests.swift
index bc6a9d8e..16aaf86c 100644
--- a/Tests/FFAITests/Ops/OpsTests.swift
+++ b/Tests/FFAITests/Ops/OpsTests.swift
@@ -268,6 +268,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 {

From aae1c601ac6d014d49b7b5fa1662a79a3d5cb700 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 07:35:58 -0500
Subject: [PATCH 08/54] =?UTF-8?q?Ops:=20add=20rmsNormRowsTwo=20=E2=80=94?=
 =?UTF-8?q?=20pair=20RMSNorm=20rows=20on=20one=20encoder?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two `mt_rms_norm_*` dispatches inside one compute encoder, used by
the Qwen3 attention mixer's pre-RoPE Q-norm + K-norm pair: both
projections run the fast path (TPG = rowSize / 4) and have no data
dependency between them. The wrapper shares a single eps buffer
allocation when the two eps values are equal — the common case
since both norms come from the same checkpoint config.

The wrapper restricts callers to the fast `mt_rms_norm_*` path
(rowSize ≤ 4096 and a multiple of 128). Wider rows must fall back
to two separate `rmsNormRows` calls so the wide kernel is picked
per call.

Test verifies bit-equivalence against two sequential `rmsNormRows`
calls on the same inputs.

Migrated from PR #10.
---
 Sources/FFAI/Ops/Ops.swift         | 93 ++++++++++++++++++++++++++++++
 Tests/FFAITests/Ops/OpsTests.swift | 61 ++++++++++++++++++++
 2 files changed, 154 insertions(+)

diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift
index 4564df7b..07c97c4e 100644
--- a/Sources/FFAI/Ops/Ops.swift
+++ b/Sources/FFAI/Ops/Ops.swift
@@ -562,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
diff --git a/Tests/FFAITests/Ops/OpsTests.swift b/Tests/FFAITests/Ops/OpsTests.swift
index 16aaf86c..62553537 100644
--- a/Tests/FFAITests/Ops/OpsTests.swift
+++ b/Tests/FFAITests/Ops/OpsTests.swift
@@ -419,6 +419,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 {

From fb10e384829edb4587dd8d368613b49a03df3f53 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 07:37:36 -0500
Subject: [PATCH 09/54] =?UTF-8?q?Ops:=20add=20kvCacheUpdateKV=20=E2=80=94?=
 =?UTF-8?q?=20append=20K=20and=20V=20on=20one=20encoder?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Wraps two `kv_cache_update_*` dispatches in a single compute
encoder. Cache layouts, dtypes, and writing thread are identical
to `kvCacheUpdate`; this only saves the encoder begin/end pair
between the K-cache write and the V-cache write. With 10 full
attention layers on Qwen3.6-A3B, this drops 10 encoder pairs per
decode token.

Migrated from PR #10.
---
 Sources/FFAI/Ops/Ops.swift                    | 47 +++++++++++++++++++
 Tests/FFAITests/Ops/OpsSpecialPathTests.swift | 30 ++++++++++++
 2 files changed, 77 insertions(+)

diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift
index 07c97c4e..80e196b1 100644
--- a/Sources/FFAI/Ops/Ops.swift
+++ b/Sources/FFAI/Ops/Ops.swift
@@ -2289,6 +2289,53 @@ 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()
+    }
+
     /// 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.
diff --git a/Tests/FFAITests/Ops/OpsSpecialPathTests.swift b/Tests/FFAITests/Ops/OpsSpecialPathTests.swift
index 2ef6836e..8575913b 100644
--- a/Tests/FFAITests/Ops/OpsSpecialPathTests.swift
+++ b/Tests/FFAITests/Ops/OpsSpecialPathTests.swift
@@ -247,6 +247,36 @@ struct OpsSpecialPathTests {
 
     // MARK: - KV cache append
 
+    @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 {

From ae534022b49015057c125812211b11a10aa831e2 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 07:39:04 -0500
Subject: [PATCH 10/54] =?UTF-8?q?Ops:=20add=20siluCastF32PlusCastF32Two=20?=
 =?UTF-8?q?=E2=80=94=20silu+cast=20plus=20two=20plain=20casts=20on=20one?=
 =?UTF-8?q?=20encoder?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Used inside the Qwen3.5 GDN T-loop where, every token, one tensor
goes through `siluCastToF32` (the gate input) while the two raw
gate scalars `aRaw` / `bRaw` 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 — all
inside one compute encoder. Saves the encoder begin/end pairs the
sequential 3-call path would create per GDN layer per token.

Migrated from PR #10.
---
 Sources/FFAI/Ops/Ops.swift                    | 61 +++++++++++++++++++
 Tests/FFAITests/Ops/OpsSpecialPathTests.swift | 44 +++++++++++++
 2 files changed, 105 insertions(+)

diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift
index 80e196b1..4fcdf7eb 100644
--- a/Sources/FFAI/Ops/Ops.swift
+++ b/Sources/FFAI/Ops/Ops.swift
@@ -4188,6 +4188,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
diff --git a/Tests/FFAITests/Ops/OpsSpecialPathTests.swift b/Tests/FFAITests/Ops/OpsSpecialPathTests.swift
index 8575913b..3be62365 100644
--- a/Tests/FFAITests/Ops/OpsSpecialPathTests.swift
+++ b/Tests/FFAITests/Ops/OpsSpecialPathTests.swift
@@ -85,6 +85,50 @@ struct OpsSpecialPathTests {
         }
     }
 
+    @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 {

From ff27159c162ca93f95d9a25b2c2d29b545ee48ee Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 07:46:07 -0500
Subject: [PATCH 11/54] =?UTF-8?q?Ops:=20GPU=20MoE=20router=20=E2=80=94=20m?=
 =?UTF-8?q?oeRouterTopK=20+=20dequantGemvInt4ExpertIndexed?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The Qwen3-MoE decode path historically had to commit + wait on a
single GPU layer just to read the router logits back to the host,
run a CPU top-K + softmax, and then dispatch the per-expert gate +
up + down GEMVs. Two new wrappers move both halves to the GPU:

- `moeRouterTopK` / `moeRouterTopKMany` — wrap the
  `mt_moe_router_topk_*` kernel. Single-row form takes
  `logits: [nExperts]` and writes `[k]` u32 indices plus `[k]`
  weights; the T-batched form iterates `program_id<0>()` across
  prefill rows in one dispatch.
- `dequantGemvInt4ExpertIndexed` /
  `dequantGemvInt4ExpertIndexedMany` — wrap
  `dequant_gemv_int4_expert_indexed_*`, which reads
  `expertIndex[0]` at kernel entry to pick which slab of a
  stacked `[nExperts, outDim, inDim/8]` u32 weight tensor to
  dequantize. The `Many` variant runs N sequential dispatches
  inside one compute encoder sharing the constexpr binding so
  Qwen3.6-A3B's topK=8 gate + up phase collapses to one encoder
  per MoE layer.

The router output indices buffer can be aliased into N single-u32
views (`buffer + offset = slot · 4`) and threaded straight into
the expert-indexed GEMV — no host readback per layer.

Tests:
- `dequantGemvInt4ExpertIndexedCorrectness` — stacks 4 experts'
  weights, picks expert 2, and checks the output matches a
  CPU-truth dequant + GEMV of that expert's slab.
- `moeRouterTopKDeterministic` — top-K of (1, 5, 3, 4) at k=2
  picks indices {1, 3}; weights sum to 1.0 and the larger logit
  dominates.

Migrated from PR #10 ITERs 54–60.
---
 Sources/FFAI/Ops/Ops.swift                  | 227 ++++++++++++++++++++
 Tests/FFAITests/Ops/QuantizedOpsTests.swift |  94 ++++++++
 2 files changed, 321 insertions(+)

diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift
index 4fcdf7eb..290fb28b 100644
--- a/Sources/FFAI/Ops/Ops.swift
+++ b/Sources/FFAI/Ops/Ops.swift
@@ -1679,6 +1679,233 @@ public enum Ops {
         enc.endEncoding()
     }
 
+    /// 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.
diff --git a/Tests/FFAITests/Ops/QuantizedOpsTests.swift b/Tests/FFAITests/Ops/QuantizedOpsTests.swift
index 4b3739d5..22d6a631 100644
--- a/Tests/FFAITests/Ops/QuantizedOpsTests.swift
+++ b/Tests/FFAITests/Ops/QuantizedOpsTests.swift
@@ -869,4 +869,98 @@ struct QuantizedOpsTests {
             }
         }
     }
+
+    @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)
+        }
+    }
 }

From 322ee15bedde829265ea29ce94c0dd2dade18b97 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 07:56:30 -0500
Subject: [PATCH 12/54] =?UTF-8?q?Ops:=20sdpaMulti=20=E2=80=94=20route=20he?=
 =?UTF-8?q?ad=5Fdim=20256=20to=20the=20d256=20kernel=20variant?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Eric's `sdpaMulti` on dev only switches on dtype — every call lands
on `ffai_sdpa_multi_*` (the head_dim=128-only kernel). Qwen3.6-A3B's
full-attention layers run at head_dim=256, and metaltile emits a
dedicated `ffai_sdpa_multi_d256_*` family with a 2-phase output
reduction that keeps the per-row threadgroup memory under Apple's
32 KB cap. Without the d256 routing the d128 kernel reads past
each row of K/V and produces silent garbage.

Switches `sdpaMulti` to a `(dtype, headDim)` tuple match so f32 /
f16 / bf16 each pick the correct kernel; relaxes
`OpsValidation.validateSdpaMulti` to accept `headDim ∈ {128, 256}`
and updates the matching validator test.

Adds a runtime test that mirrors the existing d=128 uniform-K case
at head_dim=256, exercising the new routing end-to-end.

Migrated from PR #10 ITER 93 (Qwen3.6-A3B prefill wiring needs it).
---
 Sources/FFAI/Ops/Ops.swift                   | 46 +++++++++++++++++---
 Sources/FFAI/Ops/OpsValidation.swift         |  5 ++-
 Tests/FFAITests/Ops/OpsTests.swift           | 45 +++++++++++++++++++
 Tests/FFAITests/Ops/OpsValidationTests.swift | 16 +++++--
 4 files changed, 101 insertions(+), 11 deletions(-)

diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift
index 290fb28b..6262bb87 100644
--- a/Sources/FFAI/Ops/Ops.swift
+++ b/Sources/FFAI/Ops/Ops.swift
@@ -2848,8 +2848,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,
@@ -2858,7 +2865,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,
@@ -2867,7 +2874,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,
@@ -2876,8 +2883,37 @@ 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)")
+            fatalError(
+                "Ops.sdpaMulti: unsupported dtype \(q.dtype) at headDim \(headDim) "
+                    + "— only (f32/f16/bf16) × (128, 256) are emitted")
         }
         return result
     }
diff --git a/Sources/FFAI/Ops/OpsValidation.swift b/Sources/FFAI/Ops/OpsValidation.swift
index 22f762b2..ca31d1de 100644
--- a/Sources/FFAI/Ops/OpsValidation.swift
+++ b/Sources/FFAI/Ops/OpsValidation.swift
@@ -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/Tests/FFAITests/Ops/OpsTests.swift b/Tests/FFAITests/Ops/OpsTests.swift
index 62553537..407fbd5d 100644
--- a/Tests/FFAITests/Ops/OpsTests.swift
+++ b/Tests/FFAITests/Ops/OpsTests.swift
@@ -1062,6 +1062,51 @@ struct OpsTests {
         }
     }
 
+    @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..ddb135f7 100644
--- a/Tests/FFAITests/Ops/OpsValidationTests.swift
+++ b/Tests/FFAITests/Ops/OpsValidationTests.swift
@@ -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")

From 90e2695aae2e2df7b16011d6ad14b266d094a946 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 08:00:55 -0500
Subject: [PATCH 13/54] Ops: add castToF32Two/Three, swigluMany,
 gatedMixerNormMany
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Eric's dev has the base versions (`castToF32`, `swiglu`,
`gatedMixerNorm`) but not the batched variants our PR9/10 wirings
rely on. Each variant reuses the same metaltile kernel and
collapses N independent dispatches onto a single compute encoder:

- `castToF32Two` / `castToF32Three` — Qwen3 GDN fused-prep path
  promotes 2-3 same-dtype tensors (convAct / aRaw / bRaw) to fp32
  every token; one shared encoder replaces 2-3 separate ones.
- `swigluMany` — MoE per-expert SwiGLU at decode runs topK
  independent dispatches per layer; one shared encoder per call.
- `gatedMixerNormMany` — Qwen3.5 GDN mixer's per-token norm is
  evaluated `T·Hv` times during prefill; widening the grid X axis
  from `Hv` to `T·Hv` collapses the T-loop into a single dispatch.

Tests in `Tests/FFAITests/Ops/OpsSpecialPathTests.swift` verify
each batched form matches the corresponding sequential reference.

Migrated from PR #10 (closes the Tier 2 paired-on-encoder gap
identified during the Eric-overlap audit).
---
 Sources/FFAI/Ops/Ops.swift                    | 222 ++++++++++++++++++
 Tests/FFAITests/Ops/OpsSpecialPathTests.swift | 182 ++++++++++++++
 2 files changed, 404 insertions(+)

diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift
index 6262bb87..c48004b0 100644
--- a/Sources/FFAI/Ops/Ops.swift
+++ b/Sources/FFAI/Ops/Ops.swift
@@ -4413,6 +4413,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
@@ -4605,6 +4702,86 @@ 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 scalar-sigmoid fan-out + FMA ───────────────────────────
     //
     // Wraps `mt_sigmoid_scalar_fma_{f32,f16,bf16}`. Computes
@@ -4803,6 +4980,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/Tests/FFAITests/Ops/OpsSpecialPathTests.swift b/Tests/FFAITests/Ops/OpsSpecialPathTests.swift
index 3be62365..a01643aa 100644
--- a/Tests/FFAITests/Ops/OpsSpecialPathTests.swift
+++ b/Tests/FFAITests/Ops/OpsSpecialPathTests.swift
@@ -85,6 +85,188 @@ 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 {

From 11df54f9e91c104e0adb0660db231d8b994920f4 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 08:06:33 -0500
Subject: [PATCH 14/54] =?UTF-8?q?Ops:=20batched-prefill=20wrappers=20?=
 =?UTF-8?q?=E2=80=94=20ropePartial(Many|ManyTwo),=20kvCacheUpdateKVMany,?=
 =?UTF-8?q?=20conv1dCausalStepSiluCastMany,=20gatedDeltaPrepChunk?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Eric's dev only has the single-token forms of these ops; the
PR9/10 forwardMany / decodeMany prefill paths replace the per-token
T-loop with a single batched dispatch.

- `ropePartialMany` — single dispatch covering T rows of `qk`
  rotated by per-row `positions[T]`; saves T-1 encoder pairs and
  matches T sequential `ropePartial` calls bit-identically.
- `ropePartialManyTwo` — Q + K share `positions` and constants on
  one encoder; row strides differ.
- `kvCacheUpdateKVMany` — writes T tokens' K + V rows in one
  encoder, indexed by `positions[T]`.
- `conv1dCausalStepSiluCastMany` — `T` depthwise conv + SiLU +
  fp32 cast collapsed into one dispatch with the conv state held
  in per-channel registers across the sweep.
- `gatedDeltaPrepChunk` — chunked-prefill counterpart to
  `gatedDeltaPrepStep`; runs the fused GDN prep + recurrence over
  a token range in one dispatch, keeping per-head recurrent state
  in registers across the chunk. `tLen` arrives as a `[1]` u32 GPU
  buffer so chunked dispatch chains avoid CPU readback.

All kernels (`ffai_rope_llama_many_*`,
`kv_cache_update_many_*`,
`ffai_conv1d_causal_step_silu_cast_many_*`,
`mt_gated_delta_prep_chunk_*`) are already emitted on the dev
metaltile.

Tests:
- `ropePartialManyMatchesSequential` — T=3 batched matches three
  sequential `ropePartial` outputs bit-identically.
- `kvCacheUpdateKVManyF32` — verifies multi-position K + V writes
  land in the right cache rows and leave untouched slots zero.

Migrated from PR #10 (ITERs 86–99).
---
 Sources/FFAI/Ops/Ops.swift                    | 371 ++++++++++++++++++
 Tests/FFAITests/Ops/OpsSpecialPathTests.swift |  43 ++
 Tests/FFAITests/Ops/OpsTests.swift            |  48 +++
 3 files changed, 462 insertions(+)

diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift
index c48004b0..21af2564 100644
--- a/Sources/FFAI/Ops/Ops.swift
+++ b/Sources/FFAI/Ops/Ops.swift
@@ -1044,6 +1044,134 @@ public enum Ops {
         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
@@ -1960,6 +2088,79 @@ public enum Ops {
         }
     }
 
+    /// 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
@@ -2563,6 +2764,66 @@ public enum Ops {
         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.
@@ -4196,6 +4457,116 @@ public enum Ops {
         }
     }
 
+    /// 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
+    ) {
+        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.gatedDeltaPrepChunk: every tensor must share dtype")
+        precondition(
+            dk % 32 == 0,
+            "Ops.gatedDeltaPrepChunk: dk (\(dk)) must be a multiple of 32")
+        precondition(
+            dv % 32 == 0,
+            "Ops.gatedDeltaPrepChunk: dv (\(dv)) must be a multiple of 32")
+        precondition(
+            hv % hk == 0,
+            "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,
+            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_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,
+                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,
+                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_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,
+                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,
+                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_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,
+                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,
+                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.gatedDeltaPrepChunk: unsupported dtype \(convOut.dtype)")
+        }
+    }
+
     // ─── Dynamic-M batched-prefill int4 qmm ───────────────────────────
     //
     // Host-side driver around `mt_qmm_mma`. The kernel's dispatch grid is
diff --git a/Tests/FFAITests/Ops/OpsSpecialPathTests.swift b/Tests/FFAITests/Ops/OpsSpecialPathTests.swift
index a01643aa..1f231248 100644
--- a/Tests/FFAITests/Ops/OpsSpecialPathTests.swift
+++ b/Tests/FFAITests/Ops/OpsSpecialPathTests.swift
@@ -473,6 +473,49 @@ struct OpsSpecialPathTests {
 
     // 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 {
diff --git a/Tests/FFAITests/Ops/OpsTests.swift b/Tests/FFAITests/Ops/OpsTests.swift
index 407fbd5d..8806f973 100644
--- a/Tests/FFAITests/Ops/OpsTests.swift
+++ b/Tests/FFAITests/Ops/OpsTests.swift
@@ -149,6 +149,54 @@ struct OpsTests {
         }
     }
 
+    @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 {

From 991c21450bf9d3218650094d861af7854ef7fe9a Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 08:08:13 -0500
Subject: [PATCH 15/54] =?UTF-8?q?Ops:=20add=20dequantGemvInt4Many=20?=
 =?UTF-8?q?=E2=80=94=20N=20projections,=20distinct=20inputs,=20one=20encod?=
 =?UTF-8?q?er?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Mirror of `dequantGemvInt4Two/Three/Four` for the per-expert MoE
down phase, where each chosen expert has its own activation tensor
but all N projections share PSO, `inDim`, and `groupSize`. The
wrapper sets the constexpr binding once and rotates
`(weight, scales, biases, input, output)` per dispatch, saving
N-1 encoder begin/end pairs versus N independent
`dequantGemvInt4` calls.

Test verifies the batched form matches three sequential
`dequantGemvInt4` CPU-truth references across distinct
(scale, bias) projections + distinct inputs.

Migrated from PR #10.
---
 Sources/FFAI/Ops/Ops.swift                  | 58 +++++++++++++++++++++
 Tests/FFAITests/Ops/QuantizedOpsTests.swift | 44 ++++++++++++++++
 2 files changed, 102 insertions(+)

diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift
index 21af2564..6e02eff0 100644
--- a/Sources/FFAI/Ops/Ops.swift
+++ b/Sources/FFAI/Ops/Ops.swift
@@ -1807,6 +1807,64 @@ public enum Ops {
         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()
+    }
+
     /// 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
diff --git a/Tests/FFAITests/Ops/QuantizedOpsTests.swift b/Tests/FFAITests/Ops/QuantizedOpsTests.swift
index 22d6a631..87818f21 100644
--- a/Tests/FFAITests/Ops/QuantizedOpsTests.swift
+++ b/Tests/FFAITests/Ops/QuantizedOpsTests.swift
@@ -870,6 +870,50 @@ struct QuantizedOpsTests {
         }
     }
 
+    @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 {

From 70013e00132610585b918e990c5f449d0c3f5bd6 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 08:14:22 -0500
Subject: [PATCH 16/54] =?UTF-8?q?Ops:=20add=20sdpaMultiTreeMask=20?=
 =?UTF-8?q?=E2=80=94=20tree-causal=20SDPA=20for=20tree-verify=20spec=20dec?=
 =?UTF-8?q?ode?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Same dispatch contract as `sdpaMulti(causal: false)`, but the
scalar `causal` flag is replaced by an additive `[nQuery, nQuery]`
mask tensor consulted only for in-block KV positions. Positions
`< baseKV` (the cached prefix) are always attended. Mask is `0.0`
to allow and `-inf` to block, added directly to 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 the drafter produces from its node structure.
Wraps `ffai_sdpa_multi_tree_mask_*` (headDim=128 only — no d256
variant for tree-verify on dev).

Tests:
- All-zero mask reproduces `sdpaMulti(causal: false)` output
  bit-equivalent (sanity check on the mask add path).
- Diagonal-only -inf mask isolates each query to its own KV slot:
  with uniform Q/K and per-row distinct V, output collapses to a
  copy of the matching V row.

Migrated from PR #10 ITER 70.
---
 Sources/FFAI/Ops/Ops.swift         |  86 ++++++++++++++++++++++++
 Tests/FFAITests/Ops/OpsTests.swift | 102 +++++++++++++++++++++++++++++
 2 files changed, 188 insertions(+)

diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift
index 6e02eff0..6dae62d1 100644
--- a/Sources/FFAI/Ops/Ops.swift
+++ b/Sources/FFAI/Ops/Ops.swift
@@ -3237,6 +3237,92 @@ public enum Ops {
         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.sdpaMultiTreeMask: unsupported dtype \(q.dtype)")
+        }
+        return result
+    }
+
     /// Multi-query **bidirectional** SDPA — every query attends the
     /// full `[0, baseKV + nQuery)` range. Specialized variants for
     /// head_dim ∈ {32, 64, 72} cover the VLM vision-tower spectrum
diff --git a/Tests/FFAITests/Ops/OpsTests.swift b/Tests/FFAITests/Ops/OpsTests.swift
index 8806f973..15c4265a 100644
--- a/Tests/FFAITests/Ops/OpsTests.swift
+++ b/Tests/FFAITests/Ops/OpsTests.swift
@@ -1110,6 +1110,108 @@ struct OpsTests {
         }
     }
 
+    @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 {

From fcd448f4d1c75b3e97adac8563db1af4fb3ff396 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 08:17:07 -0500
Subject: [PATCH 17/54] =?UTF-8?q?Ops:=20add=20sdpaDecode2Pass=20=E2=80=94?=
 =?UTF-8?q?=20Flash-Decoding=20two-pass=20SDPA?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Pass 1 partitions the KV cache into `blocks` even tiles and
writes per-block partial `(O, max, lse)` triples; pass 2 merges
them with the log-sum-exp trick. Total work matches single-pass
`sdpaDecode` but the two passes parallelize across more GPU
threadgroups, paying off at long KV (≥ 1K typical) where the
single-pass kernel starves the GPU with one TG per Q-head.

Block-count contract (the load-bearing detail PR10's port
silently violated): the pass2 kernel hardcodes a 32-lane block
reduction (`bn = 32`, `block_chunks = blocks / bn`), so `blocks`
MUST be a multiple of 32 AND at least 32. Anything smaller makes
`block_chunks` round to zero — pass2's outer loop iterates zero
times and the output stays whatever pass1 left in scratch,
typically all zeros. The previous port's "power of 2" check let
`blocks=2` through and was the cause of the all-zero smoke
failures during the earlier port attempt; the new precondition
catches it at the wrapper.

Scratch buffers (`partialO`, `partialM`, `partialL`) are
caller-owned so a decode loop can amortise them across steps.

Test: `blocks=32, nKV=64, nQHeads=2` — verifies the two-pass
output matches single-pass `sdpaDecode` to within 1e-2 absolute.

Migrated from PR #10 ITER 53.
---
 Sources/FFAI/Ops/Ops.swift         | 132 +++++++++++++++++++++++++++++
 Tests/FFAITests/Ops/OpsTests.swift |  69 +++++++++++++++
 2 files changed, 201 insertions(+)

diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift
index 6dae62d1..19dd1d03 100644
--- a/Sources/FFAI/Ops/Ops.swift
+++ b/Sources/FFAI/Ops/Ops.swift
@@ -3127,6 +3127,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,
diff --git a/Tests/FFAITests/Ops/OpsTests.swift b/Tests/FFAITests/Ops/OpsTests.swift
index 15c4265a..cc30ca46 100644
--- a/Tests/FFAITests/Ops/OpsTests.swift
+++ b/Tests/FFAITests/Ops/OpsTests.swift
@@ -1110,6 +1110,75 @@ 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 {

From d8183758bb8cb4df5d41bef38f6197fe1d1dfdd5 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 08:39:52 -0500
Subject: [PATCH 18/54] Ops: add batchedQkvQgemvInt4Fast + batchedQkvQmmFast
 (fused Q+K+V projection)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Wraps `ffai_batched_qkv_qgemv_fast` (M=1 decode) and
`ffai_batched_qkv_qmm_fast` (M>1 batched prefill). Both fuse the
three Q / K / V int4 dequant-projection dispatches into ONE kernel
launch — `program_id<2>()` selects the matrix across a depth-3
grid, so all three projections share TG memory and read the input
activation once instead of three times from DRAM.

Decode path (M=1): replaces the 3-dispatch `dequantGemvInt4Three`
shared-encoder form. Saves 2 encoder begin/end pairs per attention
layer.

Prefill path (M>1): outputs land in three separate contiguous
tensors (`qBuf [M, out_q]`, `kBuf [M, out_k]`, `vBuf [M, out_v]`)
so downstream code reads each projection as `[M, dim]` without
strided views. Pre-zeroes the outputs since the kernel's tile
mechanic relies on smaller-matrix tile-past-out_* slots being
already-zero.

Kernel constraints (preconditioned): `in_dim % 512 == 0`, each
`out_q / out_k / out_v % 8 == 0`, `group_size = 64`, TPG = 64.
Dispatches f32 / f16 / bf16 variants.
---
 Sources/FFAI/Ops/Ops.swift | 228 +++++++++++++++++++++++++++++++++++++
 1 file changed, 228 insertions(+)

diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift
index 19dd1d03..fde48746 100644
--- a/Sources/FFAI/Ops/Ops.swift
+++ b/Sources/FFAI/Ops/Ops.swift
@@ -1865,6 +1865,234 @@ public enum Ops {
         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
+        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.ffai_batched_qkv_qgemv_fast_f32(
+                x: x.buffer, xOffset: x.offset,
+                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.ffai_batched_qkv_qgemv_fast_f16(
+                x: x.buffer, xOffset: x.offset,
+                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.ffai_batched_qkv_qgemv_fast_bf16(
+                x: x.buffer, xOffset: x.offset,
+                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.batchedQkvQgemvInt4Fast: unsupported dtype \(x.dtype)")
+        }
+    }
+
+    /// 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.
+    ///
+    /// 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(
+            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
+        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.ffai_batched_qkv_qmm_fast_f32(
+                x: x.buffer, xOffset: x.offset,
+                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.ffai_batched_qkv_qmm_fast_f16(
+                x: x.buffer, xOffset: x.offset,
+                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.ffai_batched_qkv_qmm_fast_bf16(
+                x: x.buffer, xOffset: x.offset,
+                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.batchedQkvQmmFast: 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

From 01f591c8451756941bfdbe9ef212bf370f5060a1 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 08:41:28 -0500
Subject: [PATCH 19/54] Ops: add batched4QgemvInt4Fast + batched4QmmFast (fused
 4-output for GDN input proj)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Wraps `ffai_batched_4_qgemv_fast` (M=1 decode) and
`ffai_batched_4_qmm_fast` (M>1 batched prefill). Extends the QKV
(3-output) pattern to 4 projections sharing the same input — used
by the Qwen3.5 GDN mixer where qkv / z / b / a all project from
the same `xNorm` and can fuse into a single kernel launch.

Decode: replaces the 4-dispatch `dequantGemvInt4Four` shared-
encoder form. Pays the input DRAM roundtrip once instead of four
times; saves 3 encoder begin/end pairs per GDN layer.

Prefill (M>1): reads `x = [M, in_dim]` once into TG memory and
produces all 4 outputs (each `[M, out_*]`) in one dispatch.
Eliminates the 3 redundant `xNormFlat` DRAM reads of the unfused
`callMany` chain.

Kernel constraints (preconditioned): `in_dim % 512 == 0`, each
`out_a/b/c/d % 8 == 0`, `group_size = 64`, TPG = 64. Dispatches
f32 / f16 / bf16 variants.
---
 Sources/FFAI/Ops/Ops.swift | 256 +++++++++++++++++++++++++++++++++++++
 1 file changed, 256 insertions(+)

diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift
index fde48746..c4ec1641 100644
--- a/Sources/FFAI/Ops/Ops.swift
+++ b/Sources/FFAI/Ops/Ops.swift
@@ -2093,6 +2093,262 @@ public enum Ops {
         }
     }
 
+    /// 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.
+    ///
+    /// 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
+    ) {
+        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
+        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
+        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

From 97b4ca9e07dc6bedebae09f74b8fbc2234a9e927 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 08:42:58 -0500
Subject: [PATCH 20/54] Ops: add rmsNormQgemvInt4Fast +
 gatedRmsNormQgemvInt4Fast (fused norm+gemv)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Wraps `ffai_rms_norm_qgemv_fast` (RMSNorm-fused) and
`ffai_gated_rms_norm_qgemv_int4_fast` (Mamba2-style 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.

`rmsNormQgemvInt4Fast` is the workhorse for the finalNorm + lmHead
boundary (one dispatch per token) and any other pre-norm → int4-qmm
pattern in transformer blocks.

`gatedRmsNormQgemvInt4Fast` is purpose-built for the Qwen3.5 GDN
mixer post-recurrence boundary: collapses
`silu(z) · w · rmsnorm(y)` + outProj qgemv into one launch.

Kernel constraints (preconditioned): `in_dim % 512 == 0`,
`out_dim % 8 == 0`, `group_size = 64`, TPG = 64. The gated variant
additionally requires `Hv·Dv ≤ 8192` (kernel stages the full input
in TG memory). Dispatches f32 / f16 / bf16 variants.
---
 Sources/FFAI/Ops/Ops.swift | 223 +++++++++++++++++++++++++++++++++++++
 1 file changed, 223 insertions(+)

diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift
index c4ec1641..0e6b2a91 100644
--- a/Sources/FFAI/Ops/Ops.swift
+++ b/Sources/FFAI/Ops/Ops.swift
@@ -5913,6 +5913,229 @@ public enum Ops {
         }
     }
 
+    // ─── 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

From 55461a69646b65a591346e94c6bb374800c85443 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 08:44:12 -0500
Subject: [PATCH 21/54] Ops: add moeDownSwigluAccumInt4Chain8 (3-stage MoE
 fusion)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Wraps `ffai_moe_down_swiglu_accum_int4_chain8`. Collapses the
GPU-router MoE path's three back-to-back per-layer dispatches:

  1. `swigluMany`           — inner[k]  = silu(gate[k]) * up[k]
  2. `dequantGemvInt4ExpertIndexedMany` (down)
                            — out[k]    = W_down[expert_idx[k]] · inner[k]
  3. `scalarFMAChain8`      — acc[i]    = Σ_k slot_weight[k] · out[k][i]

into ONE kernel launch. 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 — 16 buffer-write-then-read pairs
per layer at topK=8.

Kernel constraints (preconditioned): `in_dim ≤ 768` (TG-memory
cap, kernel statically allocates 768 floats), `group_size = 64`,
TPG = 128. Lives in OpsFused.swift next to `scalarFMAChain8`
(the third stage it collapses). Dispatches f32 / f16 / bf16
variants.

Note: kernel ships in metaltile PRs #196–#201 (regenerated
locally into Sources/MetalTileSwift/Generated/MetalTileKernels.swift
for this branch).
---
 Sources/FFAI/Ops/OpsFused.swift | 154 ++++++++++++++++++++++++++++++++
 1 file changed, 154 insertions(+)

diff --git a/Sources/FFAI/Ops/OpsFused.swift b/Sources/FFAI/Ops/OpsFused.swift
index 2aa5f3b4..808bc5cf 100644
--- a/Sources/FFAI/Ops/OpsFused.swift
+++ b/Sources/FFAI/Ops/OpsFused.swift
@@ -542,6 +542,160 @@ extension Ops {
         }
     }
 
+    /// 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

From 141f6145ab631d07bca360c5bc8ec85aa86a469a Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 08:54:25 -0500
Subject: [PATCH 22/54] KVCache: add appendRangeOnGPUMany +
 reserveSlotsManyOnHost
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Batched-T KV append wrapper around Ops.kvCacheUpdateKVMany — replaces
the T-loop of Ops.kvCacheUpdateKV calls with one shared encoder + 2
dispatches. Plus the host-side slot reservation helper that fills a
u32 positions tensor under lengthLock. Pre-req for ITER 98 wiring
into Qwen35AttentionMixer.forwardMany.
---
 Sources/FFAI/KVCache/KVCache.swift | 43 ++++++++++++++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/Sources/FFAI/KVCache/KVCache.swift b/Sources/FFAI/KVCache/KVCache.swift
index 97a81a3e..505c93be 100644
--- a/Sources/FFAI/KVCache/KVCache.swift
+++ b/Sources/FFAI/KVCache/KVCache.swift
@@ -284,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

From 798915afcc04f983934006dd8908e793f4abe782 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 09:00:26 -0500
Subject: [PATCH 23/54] Qwen35AttentionMixer.forward: wire batched QKV + 2-pass
 SDPA + cached scratches

PR10 ITERs ported into the single-token attention decode path:

- ITER 22/28: cache sliceHeadHalves35 row-index tensors at init,
  pin to residency set (was rebuilt + memcpy'd per token).
- ITER 24: batch q+k+v projections through Ops.dequantGemvInt4Three
  when all 3 projections are QuantizedLinear int4 with matching groupSize.
- ITER 78: single-dispatch Ops.batchedQkvQgemvInt4Fast when constraints
  match (in_dim%512==0, out_*%8==0, groupSize==64). Backing scratch is
  one contiguous [qDim+kDim+vDim] buffer; gate via FFAI_NO_FUSED_QKV.
- ITER 65: Ops.gatherTwo replaces the two separate sliceHeadHalves35
  gathers with one shared-encoder dispatch.
- ITER 12/35: Ops.rmsNormRowsTwo collapses q_norm + k_norm into one
  shared encoder, cached scratch outputs.
- ITER 21: Ops.ropePartialTwo pairs Q+K RoPE on one shared encoder.
- ITER 53: Ops.sdpaDecode2Pass (Flash-Decoding) kicks in at long KV
  (default threshold 1024, FFAI_SDPA_2PASS_THRESHOLD overrides).
- ITER 11/40/43: Ops.sigmoidMul with cached output scratch replaces
  the mul+sigmoid chain at the attention output gate.
---
 Sources/FFAI/Models/Text/Qwen3xText.swift | 296 +++++++++++++++++++---
 1 file changed, 266 insertions(+), 30 deletions(-)

diff --git a/Sources/FFAI/Models/Text/Qwen3xText.swift b/Sources/FFAI/Models/Text/Qwen3xText.swift
index bc9a46c6..007fe2b9 100644
--- a/Sources/FFAI/Models/Text/Qwen3xText.swift
+++ b/Sources/FFAI/Models/Text/Qwen3xText.swift
@@ -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,
@@ -1499,12 +1499,64 @@ public final class Qwen35AttentionMixer: Module {
     let ropeTheta: Float
     let attnOutputGate: Bool
     let scale: Float
+    /// ITER 22 (PR10): 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, ...
+    /// ITER 24 (PR10): batched QKV-projection fast-path detection. When
+    /// all 3 projections (q, k, v) are QuantizedLinear int4 with 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?
+    /// ITER 78 (PR10): 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?
+    /// ITER 35 (PR10): cached q_norm + k_norm outputs (was fresh per call).
+    private var qNormedScratch: Tensor?
+    private var kNormedScratch: Tensor?
+    /// ITER 40 (PR10): cached sliceHeadHalves35 gather outputs.
+    private var queriesScratch: Tensor?
+    private var gateSliceScratch: Tensor?
+    /// ITER 43 (PR10): cached sdpaDecode output + sigmoidMul output.
+    private var attnOutScratch: Tensor?
+    private var attnFlatGatedScratch: Tensor?
+    /// ITER 53 (PR10): 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 +1571,61 @@ public final class Qwen35AttentionMixer: Module {
         self.ropeTheta = ropeTheta
         self.attnOutputGate = attnOutputGate
         self.scale = 1.0 / Float(Double(headDim).squareRoot())
+        if attnOutputGate {
+            // ITER 22: 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)
+            // ITER 28: pin idx scratches to the residency set.
+            device.markWeightsResident([firstBuf, secondBuf])
+        } else {
+            self.sliceFirstIdx = nil
+            self.sliceSecondIdx = nil
+        }
+        // ITER 24: 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
+            )
+            // ITER 78: 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 +1649,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)
+        // ITER 24 (PR10): batch q+k+v projections into one encoder when
+        // all 3 are QuantizedLinear int4. ITER 78: when 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 {
+                // ITER 78: 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 (ITER 24).
+                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.
+            // ITER 22: 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)
+            }
+            // ITER 65: 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)
+        // ITER 12 + ITER 35: 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. ITER 21: q + k on shared encoder.
+        Ops.ropePartialTwo(
+            qNormed, kNormed, position: position,
             headDim: headDim, rotaryDim: rotaryDim,
             thetaBase: ropeTheta, on: cmd)
 
@@ -1587,16 +1779,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)
+        // ITER 43: 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)
+        }
+        // ITER 53: 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.
+        // ITER 11 + 40: 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)
     }

From 13aa4a95b7c771a300086e2b0e1ee1fa18f9a1bd Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 09:00:53 -0500
Subject: [PATCH 24/54] Generation: add Drafter protocol + NGramDrafter +
 tree-verify scaffolding
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Speculative-decode candidate-proposer interface used by the
in-progress `SpecDecode` driver. Three implementations land here:

- `NGramDrafter` — prompt-lookup n-gram. Scans `history` for the
  last `nMatch` tokens, falls back to shorter n-grams down to
  `minNMatch`, returns up to `gamma` follow-on tokens from the
  earlier occurrence. Zero ML cost; the workhorse for code and
  structured-chat workloads where prior context recurs.
- `NeverDrafter` — stub that always returns `[]`. Used to bench
  pure spec-decode driver overhead vs raw single-token decode.
- `NGramTreeDrafter` — real branching tree drafter. Conforms to
  both `Drafter` (linear γ fallback via `propose`) and
  `TreeDrafter` (`proposeTree` returning top-K continuations per
  depth) for tree-verify spec decode against
  `Ops.sdpaMultiTreeMask`.

`DraftTreeNode` is the pure-Swift tree representation, with
`flatten()`, `verify(oracleAtHistoryEnd:oracle:)`, and
`treeCausalMask()` helpers — all pure functions, no model / cache
dependencies, ready to wire into the SpecDecode driver once
`forwardManyAllLogits` lands.

`LinearTreeAdapter` lets any linear `Drafter` masquerade as a
`TreeDrafter` (degenerate single-branch tree) so the tree-verify
driver can A/B against the linear baseline using the same
`NGramDrafter` instance.

Tests in `Tests/FFAITests/Generation/DrafterTests.swift` cover
all three drafters end-to-end: 20 host-side tests, no GPU
dependencies. All pass.

Migrated from PR #10 (ITERs 61, 68, 69, 71).
---
 Sources/FFAI/Generation/Drafter.swift         | 471 ++++++++++++++++++
 Tests/FFAITests/Generation/DrafterTests.swift | 387 ++++++++++++++
 2 files changed, 858 insertions(+)
 create mode 100644 Sources/FFAI/Generation/Drafter.swift
 create mode 100644 Tests/FFAITests/Generation/DrafterTests.swift

diff --git a/Sources/FFAI/Generation/Drafter.swift b/Sources/FFAI/Generation/Drafter.swift
new file mode 100644
index 00000000..a3165b1a
--- /dev/null
+++ b/Sources/FFAI/Generation/Drafter.swift
@@ -0,0 +1,471 @@
+// Copyright 2026 Eric Kryski (@ekryski)
+//
+// 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/Tests/FFAITests/Generation/DrafterTests.swift b/Tests/FFAITests/Generation/DrafterTests.swift
new file mode 100644
index 00000000..5aa380c0
--- /dev/null
+++ b/Tests/FFAITests/Generation/DrafterTests.swift
@@ -0,0 +1,387 @@
+// 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<Int> = [0]
+        let allow1: Set<Int> = [0, 1]
+        let allow2: Set<Int> = [0, 1, 2]
+        let allow3: Set<Int> = [0, 1, 3]  // n4 — NOT n3 (sibling)
+        let allow4: Set<Int> = [0, 4]  // n2 — NOT n1 / n3 / n4
+        let allow5: Set<Int> = [0, 4, 5]  // n5 — NOT n1 / n3 / n4
+        let allowed: [Set<Int>] = [
+            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)
+    }
+}

From eddc74eabddcca91d633e0678da783ed3a1a063c Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 09:02:38 -0500
Subject: [PATCH 25/54] Qwen35AttentionMixer.forwardMany: wire batched QKV qmm
 + sdpaMulti + many-RoPE/KV
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

PR10 ITERs ported into the T-batched prefill path:

- ITER 88: Ops.batchedQkvQmmFast fuses Q/K/V projections into one
  dispatch when fusedQKVEligible (same constraints as ITER 78
  single-token path).
- ITER 74: sliceHeadHalvesManyBoth35 helper — one shared-encoder
  Ops.gatherTwo dispatch for the queries+gate split (was two
  back-to-back sliceHeadHalvesMany35 calls).
- ITER 83: Ops.rmsNormRowsTwo collapses Q/K rmsNormRows into one
  shared encoder over (T·nHeads) + (T·nKVHeads) rows.
- ITER 84+97: Ops.ropePartialManyTwo replaces the T-loop of per-row
  ropePartial calls with ONE shared-encoder dispatch over all T
  tokens (Q+K paired). Saves ~5100 dispatches per layer at T=512.
- ITER 98: kv.appendRangeOnGPUMany + reserveSlotsManyOnHost — one
  shared encoder + 2 dispatches for the K/V cache append over all T
  tokens (was a T-loop of per-token kvCacheUpdateKV calls).
- ITER 90/93: Ops.sdpaMulti collapses T-token causal attention into
  ONE dispatch. headDim=128 routes to ffai_sdpa_multi; headDim=256
  (Qwen3.6-A3B) routes to ffai_sdpa_multi_d256. Other head_dims fall
  back to the T-loop sdpaDecode form.
- ITER 14: Ops.sigmoidMul on the prefill output-gate path.
---
 Sources/FFAI/Models/Text/Qwen3xText.swift | 252 +++++++++++++++-------
 1 file changed, 174 insertions(+), 78 deletions(-)

diff --git a/Sources/FFAI/Models/Text/Qwen3xText.swift b/Sources/FFAI/Models/Text/Qwen3xText.swift
index 007fe2b9..a41e8e80 100644
--- a/Sources/FFAI/Models/Text/Qwen3xText.swift
+++ b/Sources/FFAI/Models/Text/Qwen3xText.swift
@@ -1868,101 +1868,148 @@ 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 ──────────
+        // ITER 88 (PR10): `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 (ITER 78).
+        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. ITER 74: both halves 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. ITER 83 (PR10): collapse the two separate `rmsNormRows`
+        // dispatches into one `rmsNormRowsTwo` shared encoder, mirroring
+        // the single-token `forward` path (ITER 12+35). 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.
+        // ITER 84: pair Q+K RoPE for each row on the SAME shared encoder
+        // via `ropePartialTwo`. ITER 97: replace the T-loop of
+        // `Ops.ropePartialTwo` calls with ONE `Ops.ropePartialManyTwo`
+        // dispatch over all T tokens. 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 — ITER 98: 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 — ITER 90 / 93: 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 — ITER 14: use fused sigmoidMul on
+        // the prefill path (same kernel as the single-token ITER 11).
         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)
@@ -2949,6 +2996,55 @@ private func sliceHeadHalves35(
     return gathered.reshaped(to: [nHeads * headDim])
 }
 
+/// T-batched both-halves variant (ITER 74). 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

From f0ceb953ad4883154c83545d850f67d6fac62a77 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 09:03:31 -0500
Subject: [PATCH 26/54] Qwen35DenseMLP + Qwen35MoEFFN shared expert: use fused
 Ops.swiglu

ITER 16/17 (PR10): replace silu+mul chain with Ops.swiglu in
Qwen35DenseMLP.forward/forwardMany and Qwen35MoEFFN.forward (single-
token) + forwardMany (batched) shared-expert paths. One fused dispatch
per call vs two (silu + mul).
---
 Sources/FFAI/Models/Text/Qwen3xText.swift | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/Sources/FFAI/Models/Text/Qwen3xText.swift b/Sources/FFAI/Models/Text/Qwen3xText.swift
index a41e8e80..f8d7398c 100644
--- a/Sources/FFAI/Models/Text/Qwen3xText.swift
+++ b/Sources/FFAI/Models/Text/Qwen3xText.swift
@@ -777,17 +777,17 @@ public final class Qwen35DenseMLP: Module {
         return out
     }
 
-    /// down(silu(gate(x)) * up(x)).
+    /// down(silu(gate(x)) * up(x)). ITER 16: use 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)
     }
 }
@@ -869,7 +869,8 @@ public final class Qwen35MoEFFN: Module {
         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)
+        // ITER 16: 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()
@@ -911,7 +912,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)
+        // ITER 16: fused SwiGLU.
+        let sharedInnerAll = Ops.swiglu(gate: sgAll, up: suAll, on: work)
         let sharedOutAll = sharedDownProj.callMany(
             sharedInnerAll, t: t,
             on: work, device: device)

From 0abc058eb40b148df64eb27b1c8a188575a678fb Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 09:04:51 -0500
Subject: [PATCH 27/54] Qwen35Model: wire fused finalNorm+lmHead via
 rmsNormQgemvInt4Fast
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

ITER 76/77/79 (PR10): collapse the finalNorm + lmHead 2-dispatch chain
into one Ops.rmsNormQgemvInt4Fast call when the lmHead is 4-bit
QuantizedLinear with constraints (in_dim%512==0, out_dim%8==0,
groupSize=64). Fires on Qwen3.6-A3B's quantized lmHead.

Wired in three forward paths via shared `qwen35FinalNormLmHead`
helper:
- Qwen35Model.forward (single-token decode) — ITER 76
- Qwen35Model._forwardManyBatched (batched-T prefill last-row) — 77
- Qwen35Model.forward(inputEmbedding:) (VLM embed-input) — ITER 79

Fallback to separate finalNorm + lmHead dispatches when constraints
not met (dense lmHead, int8/int3 quant, etc.).
---
 Sources/FFAI/Models/Text/Qwen3xText.swift | 47 ++++++++++++++++++++---
 1 file changed, 41 insertions(+), 6 deletions(-)

diff --git a/Sources/FFAI/Models/Text/Qwen3xText.swift b/Sources/FFAI/Models/Text/Qwen3xText.swift
index f8d7398c..900a7861 100644
--- a/Sources/FFAI/Models/Text/Qwen3xText.swift
+++ b/Sources/FFAI/Models/Text/Qwen3xText.swift
@@ -2318,6 +2318,36 @@ public final class Qwen35AttentionLayer: Module, DecoderLayer {
     }
 }
 
+// ─── ITER 76 (PR10): 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.
@@ -2590,8 +2620,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)
+        // ITER 76 (PR10): fused rmsNorm+qgemv when lmHead is 4-bit
+        // QuantizedLinear with constraints; falls back otherwise.
+        return qwen35FinalNormLmHead(
+            h: h, finalNorm: finalNorm, lmHead: lmHead, on: cmd)
     }
 
     /// LanguageModel-protocol entry point for chunked prefill. Delegates
@@ -2794,12 +2826,13 @@ public final class Qwen35Model: LanguageModel {
         }
 
         // ── Final norm + lm_head on the LAST row only ────────────────────
+        // ITER 77 (PR10): fused via qwen35FinalNormLmHead.
         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)
+        return qwen35FinalNormLmHead(
+            h: lastRow, finalNorm: finalNorm, lmHead: lmHead, on: cmd)
     }
 
     // ─── VLM embedding-input path ────────────────────────────────────
@@ -2849,8 +2882,10 @@ public final class Qwen35Model: LanguageModel {
             workCmd.waitUntilCompleted()
         }
 
-        let normed = finalNorm(h, on: cmd)
-        return lmHead(normed, on: cmd)
+        // ITER 79 (PR10): 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.

From 9439e681c5ca258b95e8ca1616d61b4ecc94509c Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 09:07:00 -0500
Subject: [PATCH 28/54] Qwen35GDNMixer.forward: wire batched 4-projection +
 fused silu+cast2
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

PR10 ITERs ported into the single-token GDN decode path:

- ITER 19: flip FFAI_GDN_FUSED_PREP default to ON (now opt-out via
  FFAI_GDN_NO_FUSED_PREP=1). Bench delta 9.33 → 42 tps at decode T=1
  on Qwen3.6-A3B (4.5× speedup).
- ITER 23: batched 4-projection input proj via Ops.dequantGemvInt4Four
  on one shared encoder when all 4 inProj are QuantizedLinear int4
  with same groupSize. Saves 3 encoder begin/end pairs per GDN layer.
- ITER 27: pin batched-input scratches + fused-path constants in the
  Device residency set (skips per-encoder revalidation).
- ITER 95: single-dispatch fast path via Ops.batched4QgemvInt4Fast
  when constraints match (in_dim%512==0, out_dim%8==0, groupSize=64).
  Gate via FFAI_NO_FUSED_GDN_4=1.
- ITER 81: Ops.siluCastF32PlusCastF32Two collapses silu(convOut)+cast
  AND aRaw/bRaw casts onto ONE shared encoder (was 3 separate
  castToF32 dispatches). Saves 1 encoder per GDN layer × 30 ≈ 30
  pairs/decode token.

Note: ITER 100 (batched_4_qmm_fast) wiring intentionally skipped per
task spec (workgroup imbalance at GDN shape).
---
 Sources/FFAI/Models/Text/Qwen3xText.swift | 149 ++++++++++++++++++++--
 1 file changed, 137 insertions(+), 12 deletions(-)

diff --git a/Sources/FFAI/Models/Text/Qwen3xText.swift b/Sources/FFAI/Models/Text/Qwen3xText.swift
index 900a7861..51619b20 100644
--- a/Sources/FFAI/Models/Text/Qwen3xText.swift
+++ b/Sources/FFAI/Models/Text/Qwen3xText.swift
@@ -1088,6 +1088,28 @@ public final class Qwen35GDNMixer: Module {
     let yF32Scratch: Tensor
     let yGatedScratch: Tensor
 
+    /// ITER 23 (PR10): 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 via the
+    /// residency set in init (ITER 27).
+    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?
+    /// ITER 95 (PR10): 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,
@@ -1135,13 +1157,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
+        // ITER 19 (PR10): 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)
@@ -1150,6 +1173,56 @@ public final class Qwen35GDNMixer: Module {
             shape: [numValueHeads, valueHeadDim],
             dtype: .f32, device: device)
         self.yGatedScratch = Tensor.empty(shape: [valueDim], dtype: dtype, device: device)
+
+        // ITER 23 (PR10): set up the batched 4-projection input fast
+        // path when all 4 inProj are QuantizedLinear int4 with same
+        // groupSize. ITER 27: pin scratches in the residency set.
+        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
+            )
+            // ITER 95: 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)
+            // ITER 27: 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)] {
@@ -1184,16 +1257,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]
+        // ITER 23 (PR10): 4-projection shared-encoder fast path when all
+        // inProj are int4 QuantizedLinear with same groupSize. Saves 3
+        // encoder begin/end pairs per GDN layer × 30 ≈ 1.5 ms/token.
+        // ITER 95: 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]
+        // ITER 81 (PR10): 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) ─────────────────
         //
@@ -1220,9 +1339,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)
+            // ITER 81 (PR10): 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,

From 3d306c68aec825cd42705b0f460f974b5880d3c1 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 09:09:06 -0500
Subject: [PATCH 29/54] Qwen35GDNMixer.forwardMany: wire chunked recurrence +
 conv1d-many
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

PR10 ITERs ported into the T-batched GDN prefill path:

- forwardManyChunked (default route): one Ops.gatedDeltaPrepChunk
  dispatch over all T tokens, state stays in per-lane registers across
  the full T-sweep. State DRAM traffic drops from T·(load+store) to
  1·(load+store). Opt out via FFAI_GDN_NO_PREP_CHUNK=1.
- ITER 87: Ops.castToF32Two collapses a_raw/b_raw f32 casts onto ONE
  shared encoder (was 2 separate dispatches).
- ITER 99: Ops.conv1dCausalStepSiluCastMany when convKernel==4 — all
  T conv1d_steps + silu+cast fuse into ONE dispatch, conv state in
  per-channel registers across the sweep. Saves T-1 dispatches per
  layer × 30 ≈ 15K dispatches at T=512. Gate via FFAI_NO_CONV1D_MANY=1.
- Ops.gatedMixerNormMany: T·Hv rows in one shared-encoder dispatch
  (was T per-row dispatches).

Legacy forwardMany (per-token-recurrence) path retained behind
FFAI_GDN_NO_PREP_CHUNK=1, also gets the ITER-87-style castToF32Two
collapse for the per-row aRaw/bRaw casts.

Bench delta from PR10 (Qwen3.6-A3B M5 Max, 5-run median):
  T=32  prefill: 96 → 189 tps (+97%)
  T=128 prefill: 136 → 295 tps (+116%)
  T=512 prefill: 203 → 254 tps (+22%) / best 252 → 328 tps (+30%)
---
 Sources/FFAI/Models/Text/Qwen3xText.swift | 177 +++++++++++++++++++---
 1 file changed, 160 insertions(+), 17 deletions(-)

diff --git a/Sources/FFAI/Models/Text/Qwen3xText.swift b/Sources/FFAI/Models/Text/Qwen3xText.swift
index 51619b20..d56708ce 100644
--- a/Sources/FFAI/Models/Text/Qwen3xText.swift
+++ b/Sources/FFAI/Models/Text/Qwen3xText.swift
@@ -1533,6 +1533,15 @@ public final class Qwen35GDNMixer: Module {
             "Qwen35GDNMixer.forwardMany: xNormFlat size \(xNormFlat.elementCount) ≠ T·hidden = \(t * hidden)"
         )
 
+        // PR10: 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
 
@@ -1542,8 +1551,21 @@ 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)
 
+        // ITER 87 (PR10) preview: 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,
@@ -1553,38 +1575,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,
@@ -1610,6 +1624,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 ITER 99's 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)
+
+        // ITER 87 (PR10): 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 ──────
+        // ITER 99 (PR10): 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) ────────
+        // ITER (PR10): `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 ───────────────

From 59595ec4b8e6b089a2de7ce1b136d2b691d68c7a Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 09:10:37 -0500
Subject: [PATCH 30/54] Qwen35MoEFFN.forward: shared-encoder batching + fused
 residual FMA
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

PR10 ITERs ported into the MoE single-token decode path:

- ITER 26: batch sharedGate + sharedUp qmms onto ONE shared encoder
  via Ops.dequantGemvInt4Two when both are int4 with matching
  groupSize.
- ITER 82: fold sharedExpertGate (an int4 qmm producing a single
  scalar) into the SAME shared encoder via Ops.dequantGemvInt4Three.
  Saves another encoder begin/end per MoE layer × 40 ≈ 40 pairs/token.
- ITER 34: cache shared-expert sg / su / gateLogit / result outputs
  at instance level (was Tensor.empty per call → leak fix).
- ITER 66: optional residual parameter — fold the post-FFN residual
  add into the final sigmoid+FMA via Ops.sigmoidScalarFMAResidual.
  Saves 1 dispatch + 1 [hidden] DRAM roundtrip per layer per decode.
  Wired through qwen35ApplyFFN.
---
 Sources/FFAI/Models/Text/Qwen3xText.swift | 143 ++++++++++++++++++----
 1 file changed, 117 insertions(+), 26 deletions(-)

diff --git a/Sources/FFAI/Models/Text/Qwen3xText.swift b/Sources/FFAI/Models/Text/Qwen3xText.swift
index d56708ce..f01be2f2 100644
--- a/Sources/FFAI/Models/Text/Qwen3xText.swift
+++ b/Sources/FFAI/Models/Text/Qwen3xText.swift
@@ -817,6 +817,16 @@ public final class Qwen35MoEFFN: Module {
     let sharedExpertGate: AnyLinear
     let hidden: Int
 
+    /// ITER 26/82 (PR10): 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?
+    /// ITER 34: cached `result` output for the GPU fan-out.
+    private var sharedResultScratch: Tensor?
+
     init(
         moe: MoELayer,
         sharedGateProj: AnyLinear, sharedUpProj: AnyLinear,
@@ -853,38 +863,122 @@ 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`.
+    ///
+    /// ITER 66 (PR10): 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)
+        // ITER 26: batch sharedGate + sharedUp qmms onto one shared
+        // encoder via dequantGemvInt4Two when both are int4.
+        // ITER 82: also fold sharedExpertGate (an int4 qmm producing a
+        // single scalar from the same `xNorm`) into the SAME shared
+        // encoder via dequantGemvInt4Three. Saves another encoder
+        // begin/end per MoE layer × 40 ≈ 40 pairs / decode token.
+        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)
+        }
         // ITER 16: 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 ITER 66 is in play) 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 {
+            // ITER 66: 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
     }
@@ -2705,17 +2799,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`. ITER 66 (PR10): fold 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)
     }
 }
 

From a0dcfe662545c03bc3c14c518bc5eed3ed161f4e Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 09:18:53 -0500
Subject: [PATCH 31/54] feat(moe): wire GPU router fast path in MoELayer.decode
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

ITER 56/58/64/80/94 (PR10) ported into the single-token MoE decode:

- ITER 56/80: GPU MoE router fast path — default ON via gpuRouterEnabled
  (opt out FFAI_MOE_GPU_ROUTER=0). Predicated on stacked-int4 experts,
  softmax-then-topK with normTopKProb, no expertBias, topK=8. Replaces
  the cmd.commit + waitUntilCompleted CPU sync with full-GPU pipeline.
  Net: 40 × waitUntilCompleted per decode token to 0.
- ITER 64: per-instance env-flag cache (gpuRouterEnabled, noDequantGate,
  forceBGEMMAtT1, noBGEMMBm64) so the decode hot path skips dictionary
  lookups.
- ITER 58: batched per-expert indexed qmms on ONE shared encoder per
  phase via Ops.dequantGemvInt4ExpertIndexedMany — 16 calls (8 gate +
  8 up) collapse to one encoder; down stays in its own encoder
  (different out_dim).
- ITER 94: fuse phase 1b (swigluMany) + phase 2 (8 down indexed-gemvs)
  + phase 3 (scalarFMAChain8) into ONE dispatch via
  Ops.moeDownSwigluAccumInt4Chain8 when moeIntermediate ≤ 768 +
  groupSize=64 (Qwen3.6-A3B). Falls back to the 3-dispatch chain
  (swigluMany + dequantGemvInt4ExpertIndexedMany + scalarFMAChain8)
  for larger configs.
- Per-instance scratches (router indices/weights, expert g/u/inner/out,
  accumulator) pre-allocated lazily on first invocation, indexed by
  slot. Mirrors the ITER 32/36/38/41 caching pattern.
---
 Sources/FFAI/Models/MoELayer.swift | 238 ++++++++++++++++++++++++++++-
 1 file changed, 237 insertions(+), 1 deletion(-)

diff --git a/Sources/FFAI/Models/MoELayer.swift b/Sources/FFAI/Models/MoELayer.swift
index ce8875ce..c65d3aa2 100644
--- a/Sources/FFAI/Models/MoELayer.swift
+++ b/Sources/FFAI/Models/MoELayer.swift
@@ -390,6 +390,34 @@ public final class MoELayer: Module, DecoderLayer {
     public let enableBGEMM: Bool
     public let useBm8Env: Bool
     public let useM1Env: Bool
+    /// ITER 64 (PR10): additional env-flag caches.
+    let noDequantGate: Bool
+    let gpuRouterEnabled: Bool
+    let forceBGEMMAtT1: Bool
+    let noBGEMMBm64: Bool
+
+    /// ITER 15/42 (PR10): 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?
+
+    /// ITER 41 (PR10): cached scratches for the per-expert weighted-sum
+    /// loop at decode T=1.
+    private var accumulatorScratch: Tensor?
+    private var topKScalarsBuf: MTLBuffer?
+    /// ITER 32 / 36 / 38 (PR10): cached per-expert g/u/inner/down outputs.
+    private var expertGScratches: [Tensor] = []
+    private var expertUScratches: [Tensor] = []
+    private var expertInnerScratches: [Tensor] = []
+    private var expertOutScratches: [Tensor] = []
+    /// ITER 56 (PR10): 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 +478,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
+        // ITER 80 (PR10): default GPU router ON now that ITER 72's
+        // end-to-end correctness pinning is in place across bf16/f16/f32.
+        // 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 +528,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. ITER 56 (PR10): GPU MoE router fast path ─────────────
+        // When `FFAI_MOE_GPU_ROUTER=1` (default ON per ITER 80) 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
@@ -1043,6 +1100,185 @@ public final class MoELayer: Module, DecoderLayer {
         let inner = Ops.swiglu(gate: g, up: u, on: cmd)
         return downProj(inner, on: cmd)
     }
+
+    /// ITER 56 (PR10): 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,
+    /// ITER 94).
+    ///
+    /// 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)
+
+        // ITER 58 (PR10): 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)
+
+        // ITER 94 (PR10): 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 (ITER 29/30/31, 38/39, ITER 47) 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.
+        cmd.commit()
+        return acc
+    }
 }
 
 // ─── Integration note for MoE-bearing families ───────────────────────

From 5c0dbeaeba3232eba3edbcb3850e89e8d0564d97 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 09:20:12 -0500
Subject: [PATCH 32/54] feat(moe): wire GPU router fast path in
 MoELayer.decodeMany + bm64 default
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

ITER 92 (PR10): T-batched GPU router via Ops.moeRouterTopKMany when
router config matches the kernel envelope (softmax-then-topK, normTopK
prob, no expertBias, k=8). 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.

ITER 73 (PR10): flip bm64_mpp default-on threshold from m_total ≥ 64
(opt-in) to m_total ≥ 1024 (default). Bench crossover sits between
m_total=512 (bm16 +19%) and m_total=1024 (bm64 +18%); the threshold
matches the refreshed Qwen3.6-A3B sweep. Opt out via
FFAI_MOE_BGEMM_NO_BM64=1 (cached on self.noBGEMMBm64).
---
 Sources/FFAI/Models/MoELayer.swift | 67 +++++++++++++++++++++++-------
 1 file changed, 52 insertions(+), 15 deletions(-)

diff --git a/Sources/FFAI/Models/MoELayer.swift b/Sources/FFAI/Models/MoELayer.swift
index c65d3aa2..36a111d4 100644
--- a/Sources/FFAI/Models/MoELayer.swift
+++ b/Sources/FFAI/Models/MoELayer.swift
@@ -681,20 +681,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 ──────────────────────
+        // ITER 92 (PR10): 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 ────────────────────────
@@ -813,9 +849,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
+        // ITER 73 (PR10): bm64_mpp default-on for mTotal ≥ 1024 (refreshed
+        // bench: 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)

From f960dfc51f045f03f510de1bc410ecce1859a3e5 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 09:20:56 -0500
Subject: [PATCH 33/54] test(qwen36): add forwardManyBench T=2048 long-context
 bench case

Extends the existing forwardManyBench family (T=32, T=128, T=512)
with a T=2048 case so the bench harness covers the long-context
prefill shape where the batched forwardMany speedup should stabilise
versus the per-token decode loop. Same per-token vs batched
median-of-5 protocol; same 2-warmup convention. Just changes the
targetT value.
---
 .../Text/Qwen36TextIntegrationTests.swift                    | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/Tests/ModelIntegrationTests/Text/Qwen36TextIntegrationTests.swift b/Tests/ModelIntegrationTests/Text/Qwen36TextIntegrationTests.swift
index e97f28ec..0f67c276 100644
--- a/Tests/ModelIntegrationTests/Text/Qwen36TextIntegrationTests.swift
+++ b/Tests/ModelIntegrationTests/Text/Qwen36TextIntegrationTests.swift
@@ -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()

From 844ed396567c1f867022f5089121ee01b27a4b5a Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 09:21:36 -0500
Subject: [PATCH 34/54] feat(qwen35): add forwardManyAllLogits for per-position
 spec-decode verify
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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 via batched
rmsNormRows + lmHead.callMany.

Refactors _forwardManyBatched to accept a returnAllLogits flag so
forwardMany (last-row) and forwardManyAllLogits (all-T) share the
same KV/GDN cache machinery. forwardMany retains the ITER 77 fused
qwen35FinalNormLmHead path for the last-row case.

Spec-decode driver pattern: drafter proposes γ candidate tokens,
target verifies via forwardManyAllLogits(prefix + candidates) to get
per-position logits, then accepts up to the first reject.
---
 Sources/FFAI/Models/Text/Qwen3xText.swift | 65 +++++++++++++++++++----
 1 file changed, 55 insertions(+), 10 deletions(-)

diff --git a/Sources/FFAI/Models/Text/Qwen3xText.swift b/Sources/FFAI/Models/Text/Qwen3xText.swift
index f01be2f2..17167d8c 100644
--- a/Sources/FFAI/Models/Text/Qwen3xText.swift
+++ b/Sources/FFAI/Models/Text/Qwen3xText.swift
@@ -3068,6 +3068,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)
     }
 
@@ -3099,11 +3124,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
@@ -3184,14 +3215,28 @@ public final class Qwen35Model: LanguageModel {
             workCmd.waitUntilCompleted()
         }
 
-        // ── Final norm + lm_head on the LAST row only ────────────────────
-        // ITER 77 (PR10): fused via qwen35FinalNormLmHead.
-        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)
+        // ── Final norm + lm_head ─────────────────────────────────────────
+        // Two output shapes (see _forwardManyBatched docstring):
+        //   * returnAllLogits == false: logits of the LAST row only.
+        //     ITER 77 (PR10) 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 ────────────────────────────────────

From d0270ac99b437292499f0fca4744dc8fcb205355 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 09:24:54 -0500
Subject: [PATCH 35/54] feat(kvcache): add snapshot/restore primitives +
 composite cache snapshot
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Spec-decode rollback needs to restore ALL per-layer caches when a
draft candidate is rejected. Add the per-cache primitives + a
composite snapshotAll/restoreAll over [any LayerCacheProtocol]:

- GDNStateCache: snapshot/restore the [Hv, Dv, Dk] f32 current state
  buffer via blit copy + setLength to rewind the position counter
  (without zeroing buffers that reset() would wipe). Cached snapshot
  tensor is reused across calls.
- ConvStateCache: snapshot/restore the [K-1, C] rolling window via
  blit copy with a similarly cached snapshot tensor.
- Qwen35GDNLayerCache: setLength helper for the composite (conv+gdn)
  position counter; conv/gdn buffer restore is delegated to the
  per-cache restore methods.
- CacheSnapshot.swift: LayerCacheSnapshot enum + snapshotAll/restoreAll
  array extension routing per cache kind. Handles KVCache via
  truncate(toLength:), Qwen35GDNLayerCache via composite restore.

Per-step cost on Qwen3.6-A3B: ~150 µs aggregate blit time (60 MiB GDN
+ 900 KB conv on unified memory at ~400 GB/s).
---
 Sources/FFAI/KVCache/CacheSnapshot.swift  | 96 +++++++++++++++++++++++
 Sources/FFAI/KVCache/ConvStateCache.swift | 55 +++++++++++++
 Sources/FFAI/KVCache/GDNStateCache.swift  | 68 ++++++++++++++++
 Sources/FFAI/Models/Text/Qwen3xText.swift | 11 +++
 4 files changed, 230 insertions(+)
 create mode 100644 Sources/FFAI/KVCache/CacheSnapshot.swift

diff --git a/Sources/FFAI/KVCache/CacheSnapshot.swift b/Sources/FFAI/KVCache/CacheSnapshot.swift
new file mode 100644
index 00000000..b16fc78d
--- /dev/null
+++ b/Sources/FFAI/KVCache/CacheSnapshot.swift
@@ -0,0 +1,96 @@
+// 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 3a78a0f6..68e56104 100644
--- a/Sources/FFAI/KVCache/ConvStateCache.swift
+++ b/Sources/FFAI/KVCache/ConvStateCache.swift
@@ -62,6 +62,61 @@ 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.
+    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 b96dca8b..59150f7b 100644
--- a/Sources/FFAI/KVCache/GDNStateCache.swift
+++ b/Sources/FFAI/KVCache/GDNStateCache.swift
@@ -134,6 +134,74 @@ 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.
+    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/Models/Text/Qwen3xText.swift b/Sources/FFAI/Models/Text/Qwen3xText.swift
index 17167d8c..68f65b74 100644
--- a/Sources/FFAI/Models/Text/Qwen3xText.swift
+++ b/Sources/FFAI/Models/Text/Qwen3xText.swift
@@ -1108,6 +1108,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 ────────────────

From d45d60953a454cd355252787979f0d170497b52c Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 09:26:00 -0500
Subject: [PATCH 36/54] feat(spec-decode): add greedy speculative-decode driver
 SpecDecode
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Greedy spec-decode loop for Qwen3.5/3.6:
  1. Drafter proposes up to γ candidates via Drafter.propose.
  2. Snapshot caches; run forwardManyAllLogits([prev, c_0..c_{γ-1}])
     to advance caches by γ+1 tokens and get per-position logits.
  3. Greedy verify: accept consecutive argmax matches up to the first
     mismatch.
  4. On partial accept (any mismatch): restoreAll caches from
     snapshot, replay accepted prefix via single-step forward, take
     next-prev from the LAST single-step (bit-identical to baseline
     greedy loop; forwardManyAllLogits drifts slightly due to
     different kernel paths). On full accept: commit all candidates +
     bonus token from trailing logits[γ].
  5. Empty proposal → fall back to single-token forward.

Stops on maxNewTokens or stop-tokens. SpecDecodeStats records
acceptanceRate, tps, fallbackSingleSteps, etc.

Greedy-only for v0 — temperature / nucleus / top-K follows the same
loop with probability-ratio rejection comparing sampled tokens against
the drafter's proposed tokens.
---
 Sources/FFAI/Generation/SpecDecode.swift | 238 +++++++++++++++++++++++
 1 file changed, 238 insertions(+)
 create mode 100644 Sources/FFAI/Generation/SpecDecode.swift

diff --git a/Sources/FFAI/Generation/SpecDecode.swift b/Sources/FFAI/Generation/SpecDecode.swift
new file mode 100644
index 00000000..3c4f880a
--- /dev/null
+++ b/Sources/FFAI/Generation/SpecDecode.swift
@@ -0,0 +1,238 @@
+// 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<Int> = [],
+        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
+}

From 8de0f9bdaeb7a8ba50e896a8ebcb34d9a8ed4414 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 09:38:23 -0500
Subject: [PATCH 37/54] fix(tests): silence MTLComputePipelineState Sendable
 error in PSOCacheTests

The concurrency test uses a TaskGroup of MTLComputePipelineState, which
under Swift 6 strict concurrency tripped 'type does not conform to the
Sendable protocol' on the CI runner (Xcode 16.4 SDK).

Adding @preconcurrency to the Metal import keeps the same test
behaviour but lets the compiler treat Metal's pre-Sendable protocols
as warnings rather than errors, matching the convention the rest of
the project uses for system-framework imports under strict
concurrency.
---
 Tests/MetalTileSwiftTests/PSOCacheTests.swift | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Tests/MetalTileSwiftTests/PSOCacheTests.swift b/Tests/MetalTileSwiftTests/PSOCacheTests.swift
index d4002b6a..973fbe22 100644
--- a/Tests/MetalTileSwiftTests/PSOCacheTests.swift
+++ b/Tests/MetalTileSwiftTests/PSOCacheTests.swift
@@ -24,7 +24,7 @@
 // lookup in `PSOCache.lookup`).
 
 import Foundation
-import Metal
+@preconcurrency import Metal
 import Testing
 
 @testable import MetalTileSwift

From f1684a1ed2177725a76cb4108fce0915ac689e81 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 09:45:36 -0500
Subject: [PATCH 38/54] test(spec-decode): port end-to-end SpecDecode greedy
 bench against Qwen3.6

Integration test from PR #9 that compares baseline greedy decode
(forward loop) vs SpecDecode.generateGreedy + NGramDrafter at
gamma in 1, 2, 4. Verifies generated sequences match baseline
exactly (greedy + greedy drafter + correct snapshot/restore is
deterministic) and reports decode-only tps so the cost of the
verify forward is visible.

Gated on the local Qwen3.6-A3B checkpoint at
/Users/tom/models/Qwen3.6-35B-A3B-4bit; skips with a print when
the path is missing.
---
 .../Text/SpecDecodeBenchTests.swift           | 174 ++++++++++++++++++
 1 file changed, 174 insertions(+)
 create mode 100644 Tests/ModelIntegrationTests/Text/SpecDecodeBenchTests.swift

diff --git a/Tests/ModelIntegrationTests/Text/SpecDecodeBenchTests.swift b/Tests/ModelIntegrationTests/Text/SpecDecodeBenchTests.swift
new file mode 100644
index 00000000..a105b74e
--- /dev/null
+++ b/Tests/ModelIntegrationTests/Text/SpecDecodeBenchTests.swift
@@ -0,0 +1,174 @@
+// Copyright 2026 Eric Kryski (@ekryski)
+//
+// 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
+}

From d016ba168178c6cd9dcd5e19cf113e4ae16c12dc Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 09:46:45 -0500
Subject: [PATCH 39/54] =?UTF-8?q?test(spec-decode):=20port=20verify-cost?=
 =?UTF-8?q?=20bench=20=E2=80=94=20forwardManyAllLogits=20vs=20single-step?=
 =?UTF-8?q?=20loop?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Measures whether spec-decode's verify step should use
forwardManyAllLogits(T) or a loop of single forward() calls.

Times 16 iterations of each path with cache snapshot/restore between
iterations (so each iter starts in identical state) and reports the
median ms / iter ratio for T=3 (gamma=2 verify) and T=2 (gamma=1
verify) shapes. The answer determines the SpecDecode driver's
optimal verify primitive once the cost crossover point is found.

Gated on the local Qwen3.6-A3B checkpoint at
/Users/tom/models/Qwen3.6-35B-A3B-4bit; skips with a print when the
path is missing.
---
 .../Text/SpecDecodeVerifyCostBench.swift      | 189 ++++++++++++++++++
 1 file changed, 189 insertions(+)
 create mode 100644 Tests/ModelIntegrationTests/Text/SpecDecodeVerifyCostBench.swift

diff --git a/Tests/ModelIntegrationTests/Text/SpecDecodeVerifyCostBench.swift b/Tests/ModelIntegrationTests/Text/SpecDecodeVerifyCostBench.swift
new file mode 100644
index 00000000..9fb95334
--- /dev/null
+++ b/Tests/ModelIntegrationTests/Text/SpecDecodeVerifyCostBench.swift
@@ -0,0 +1,189 @@
+// Copyright 2026 Eric Kryski (@ekryski)
+//
+// 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.completed()
+        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.completed()
+            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.completed()
+        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.completed()
+            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))")
+    }
+}

From 2529b1c2aa7e375b596820f45614778642b9fae0 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 09:47:38 -0500
Subject: [PATCH 40/54] fix(firered-vad): break up BINUNICODE8 length
 expression so Swift can type-check it

CI build was failing on Swift 6 / Xcode 16.4 with 'the compiler is
unable to type-check this expression in reasonable time' on the
8-byte little-endian length decode at line 1202. The pre-existing
single expression chained 8 byte loads + 7 bit-shifts + 7 bit-ORs
through Int promotion, which trips the type-inference budget under
strict concurrency.

Split into two 32-bit halves (nLo32 + nHi32) joined by one final OR;
identical bit semantics, type-check finishes in &lt;1 ms.
---
 Sources/FFAI/Models/Audio/VAD/FireRedVAD.swift | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/Sources/FFAI/Models/Audio/VAD/FireRedVAD.swift b/Sources/FFAI/Models/Audio/VAD/FireRedVAD.swift
index c8c5d7a9..de623599 100644
--- a/Sources/FFAI/Models/Audio/VAD/FireRedVAD.swift
+++ b/Sources/FFAI/Models/Audio/VAD/FireRedVAD.swift
@@ -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

From 04669460366c98417e0496ed870ab7ddf36be55f Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 10:00:11 -0500
Subject: [PATCH 41/54] test(qwen35-moe): add Qwen3.5-35B-A3B forwardMany +
 decode T=1 bench harness
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Same shape as Qwen3.6's forwardManyBench but pointed at the locally-
available Qwen3.5-35B-A3B-4bit checkpoint. Engine is m.qwen35 in
both cases, so this exercises every Bagel-migration wiring landed
on tom/bagel-clean (GPU MoE router, batched QKV, GDN chunked
recurrence, rmsNormQgemvInt4Fast finalNorm, 2-pass Flash-Decoding,
MTLResidencySet weight pinning).

Each test gates on the local checkpoint existing and prints a skip
message when missing.

M5 Max decode T=1 measurement: 96.87 tps (median of 5 × 32-step
runs), matching the PR #10 Qwen3.6-A3B tip target of 96.3 tps.
---
 .../Text/Qwen35MoEBenchIntegrationTests.swift | 178 ++++++++++++++++++
 1 file changed, 178 insertions(+)
 create mode 100644 Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift

diff --git a/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift b/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift
new file mode 100644
index 00000000..2fe7d222
--- /dev/null
+++ b/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift
@@ -0,0 +1,178 @@
+// Copyright 2026 Eric Kryski (@ekryski)
+//
+// 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 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)
+    }
+
+    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.completed()
+        }
+
+        // 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.completed()
+            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))")
+    }
+}

From d42df3b2e2e66eabc6fe81800af7bf663637a471 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 10:07:37 -0500
Subject: [PATCH 42/54] fix(layers): QuantizedLinear.callMany 4-bit fast path
 needs outDim % 32 guard
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

QuantizedLinear.callMany routes 4-bit weights to Ops.dequantGemmDynamicM
(mt_qmm_mma backed), whose kernel requires outDim to be a multiple of
32 (the BN tile of the MMA kernel). Without an outDim guard the
precondition tripped on Qwen3.5-35B-A3B's shared_expert_gate, which
quantises to 4-bit at outDim=1 (a per-token gate scalar):
'Ops.dequantGemmDynamicM: nOut (1) must be multiple of 32 (BN tile)'.

Add the multiple-of-32 check next to the 4-bit branch. Anything that
doesn't fit (4-bit narrow projections + every non-4-bit) falls through
to the per-row dequantGemv loop, bit-identical to the per-token path
and irrelevant to perf at outDim=1.

Verified on M5 Max with Qwen3.5-35B-A3B-4bit: forwardManyBench T=8 /
T=128 / T=512 all complete cleanly. Headline perf: decode T=1 96.87
tps, prefill T=512 458.6 tps (4.75x speedup over per-token loop) —
matches the PR #10 tip targets of 96.3 / 455 tps.
---
 Sources/FFAI/Layers.swift | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/Sources/FFAI/Layers.swift b/Sources/FFAI/Layers.swift
index 7873e3b2..5ad8a750 100644
--- a/Sources/FFAI/Layers.swift
+++ b/Sources/FFAI/Layers.swift
@@ -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,

From d80350add3c25e9a79f3ed98d3fec15f17f0b4af Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 10:25:26 -0500
Subject: [PATCH 43/54] =?UTF-8?q?perf:=20bump=20default=20maxCommandBuffer?=
 =?UTF-8?q?Count=2016=E2=86=9264=20+=20Ops.swiglu=20in=20Granite4/LFM2/Jam?=
 =?UTF-8?q?ba=20dense=20MLP?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two ITERs from PR #10 the wirings agent missed in the first pass.

ITER 20: bump MetalTileLibrary.defaultMaxCommandBufferCount from 16
to 64. Lets the Metal driver pipeline more cmd-buffers in flight
before applying backpressure; absorbs per-cmd-buffer encode latency
across more concurrent submissions. PR #10 measured +2-3% prefill
T=512 on Qwen3.6-A3B / M5 Max. FFAI_MAX_COMMAND_BUFFERS override
preserved.

ITER 17: replace Ops.mul(Ops.silu(g, cmd), u, cmd) with the fused
Ops.swiglu in Granite4Text.MLP.forward, LFM2Text.MLP.forward, and
JambaText.MLP.forward. Same semantic; one kernel dispatch instead of
two, intermediate silu(g) stays in registers instead of round-
tripping to DRAM. Brings these three cross-family models in line
with Qwen3.5 (already wired via the agent's earlier ITER 16 commit
on Qwen35DenseMLP + MoEFFN shared expert).
---
 Sources/FFAI/Models/Text/Granite4Text.swift   | 4 +++-
 Sources/FFAI/Models/Text/JambaText.swift      | 4 +++-
 Sources/FFAI/Models/Text/LFM2Text.swift       | 4 +++-
 Sources/MetalTileSwift/MetalTileLibrary.swift | 8 +++++++-
 4 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/Sources/FFAI/Models/Text/Granite4Text.swift b/Sources/FFAI/Models/Text/Granite4Text.swift
index 5bda00be..adcb27e3 100644
--- a/Sources/FFAI/Models/Text/Granite4Text.swift
+++ b/Sources/FFAI/Models/Text/Granite4Text.swift
@@ -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..fd609d97 100644
--- a/Sources/FFAI/Models/Text/JambaText.swift
+++ b/Sources/FFAI/Models/Text/JambaText.swift
@@ -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..66517bbc 100644
--- a/Sources/FFAI/Models/Text/LFM2Text.swift
+++ b/Sources/FFAI/Models/Text/LFM2Text.swift
@@ -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/MetalTileSwift/MetalTileLibrary.swift b/Sources/MetalTileSwift/MetalTileLibrary.swift
index c9af5f59..30184ebe 100644
--- a/Sources/MetalTileSwift/MetalTileLibrary.swift
+++ b/Sources/MetalTileSwift/MetalTileLibrary.swift
@@ -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. PR10 ITER 20 measured the bump
+        // 16 → 64 as a 2-3% prefill T=512 win on Qwen3.6-A3B / M5 Max
+        // by absorbing the 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

From f56a799c8b3e0b850a2fda4723260926203bec58 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 10:34:08 -0500
Subject: [PATCH 44/54] test(qwen35-moe): add T=2048 long-context +
 decode-after-1024-prefill bench cases

Extends the bench harness with two long-context shapes that PR #10's
work needs to validate but Eric's existing harness only covered up to
T=512:

- forwardManyT2K (T=2048): exercises chunked GDN recurrence over a
  longer window + larger SDPA prefill grid.
- decodeAfterLongPrefill: prefill 1024 tokens via forwardMany, then
  measure decode-only tps over 32 steps. Catches any KV-cache /
  residency / sdpaDecode2Pass routing regression at non-trivial KV.

Measured on M5 Max, Qwen3.5-35B-A3B-4bit:

  forwardManyBench T=2048: per_token=21499ms batched=4755ms
                           speedup=4.52x batched_tps=430.7

  decode T=1 after T=1024 prefill:
    runs=[0.335, 0.335, 0.336, 0.336, 0.337]s median=0.336s -> 95.37 tps

Decode tps barely drops from empty-cache (96.51) to KV=1024 (95.37)
- the kernel-compute floor PR #10 documented at 96 tps holds through
the long-context regime.
---
 .../Text/Qwen35MoEBenchIntegrationTests.swift | 111 ++++++++++++++++++
 1 file changed, 111 insertions(+)

diff --git a/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift b/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift
index 2fe7d222..3a54d5ac 100644
--- a/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift
+++ b/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift
@@ -84,6 +84,54 @@ struct Qwen35MoEBenchIntegrationTests {
         )
     }
 
+    @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.completed()
+        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)
@@ -94,6 +142,69 @@ struct Qwen35MoEBenchIntegrationTests {
         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 decode after T=1024 prefill — long-KV decode tps")
+    func decodeAfterLongPrefill() async throws {
+        guard FileManager.default.fileExists(atPath: qwen35MoELocalPath) else {
+            print("decodeAfterLongPrefill 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 < 1024 { encoded.append(contentsOf: seedEncoded) }
+        encoded = Array(encoded.prefix(1024))
+
+        // Warm.
+        for _ in 0 ..< 2 {
+            let warmCaches = qwen.makeLayerCaches()
+            let warmCmd = Device.shared.makeCommandBuffer()
+            _ = qwen.forwardMany(
+                tokenIds: encoded, startPosition: 0,
+                caches: warmCaches, on: warmCmd, device: Device.shared)
+            warmCmd.commit()
+            await warmCmd.completed()
+        }
+
+        // Prefill 1024 via forwardMany, then decode 32 steps and
+        // measure decode-only tps. Verifies the wirings (sdpaDecode +
+        // sdpaDecode2Pass routing, KV cache residency, etc.) hold up
+        // at non-trivial KV.
+        let nSteps = 32
+        var runs: [Double] = []
+        for _ in 0 ..< 5 {
+            let caches = qwen.makeLayerCaches()
+            let cmd = Device.shared.makeCommandBuffer()
+            _ = qwen.forwardMany(
+                tokenIds: encoded, startPosition: 0,
+                caches: caches, on: cmd, device: Device.shared)
+            cmd.commit()
+            await cmd.completed()
+
+            let t0 = Date()
+            for j in 0 ..< nSteps {
+                _ = qwen.forward(
+                    tokenId: 0, position: 1024 + 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 after T=1024 prefill: "
+                + "runs=\(runs.map { String(format: "%.3f", $0) })s "
+                + "median=\(String(format: "%.3f", median))s → "
+                + "\(String(format: "%.2f", tps)) tps")
+    }
+
     private func loadModel() async throws -> Model {
         var optsBuilder = LoadOptions()
         optsBuilder.prewarm = false

From ba1b2b6a35c8eef6fb47d96f1ab3d1deb91dd7de Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 10:51:49 -0500
Subject: [PATCH 45/54] test(qwen35-moe): replace T=1024 decode-after-prefill
 with T=32K long-context bench
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Drops the lukewarm 1024-token case (KV size barely above empty-cache,
decode tps barely budged from 96.51 to 95.37) and replaces it with
a real long-context probe at T=32768.

Single-run shape — 5-run median is impractical at 32K (each run is
multi-minute). Per-token baseline is intentionally not measured at
that size (would be 5+ minutes per run). Measures:

  - Batched forwardMany at T=32768.
  - Decode-only tps over 32 steps after the T=32768 prefill, so any
    KV-cache scaling regression (sdpaDecode2Pass routing, residency
    pinning, etc.) surfaces.

Measured on M5 Max / Qwen3.5-35B-A3B-4bit (single run):
  Prefill T=32768: 243.829s -> 134.4 tps batched
  Decode T=1 after T=32768 prefill: 2.215s over 32 steps -> 14.45 tps

The decode drop from 96 tps (empty KV) to 14.45 tps (32K KV) is the
expected 10-attention-layer KV-bandwidth tax — 30 of 40 layers are
GDN with O(1) state, but the 10 full-attention layers must walk the
linearly-growing K/V caches every decode step.
---
 .../Text/Qwen35MoEBenchIntegrationTests.swift | 81 +++++++++----------
 1 file changed, 39 insertions(+), 42 deletions(-)

diff --git a/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift b/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift
index 3a54d5ac..3c861f8f 100644
--- a/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift
+++ b/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift
@@ -147,10 +147,10 @@ struct Qwen35MoEBenchIntegrationTests {
         try await runForwardManyBench(targetT: 2048)
     }
 
-    @Test("Qwen3.5-35B-A3B decode after T=1024 prefill — long-KV decode tps")
-    func decodeAfterLongPrefill() async throws {
+    @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("decodeAfterLongPrefill skipped: \(qwen35MoELocalPath) not found")
+            print("longContext32K skipped: \(qwen35MoELocalPath) not found")
             return
         }
         let m = try await loadModel()
@@ -158,51 +158,48 @@ struct Qwen35MoEBenchIntegrationTests {
         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 < 1024 { encoded.append(contentsOf: seedEncoded) }
-        encoded = Array(encoded.prefix(1024))
-
-        // Warm.
-        for _ in 0 ..< 2 {
-            let warmCaches = qwen.makeLayerCaches()
-            let warmCmd = Device.shared.makeCommandBuffer()
-            _ = qwen.forwardMany(
-                tokenIds: encoded, startPosition: 0,
-                caches: warmCaches, on: warmCmd, device: Device.shared)
-            warmCmd.commit()
-            await warmCmd.completed()
+        while encoded.count < targetT {
+            encoded.append(contentsOf: seedEncoded)
         }
+        encoded = Array(encoded.prefix(targetT))
+        let T = encoded.count
 
-        // Prefill 1024 via forwardMany, then decode 32 steps and
-        // measure decode-only tps. Verifies the wirings (sdpaDecode +
-        // sdpaDecode2Pass routing, KV cache residency, etc.) hold up
-        // at non-trivial KV.
-        let nSteps = 32
-        var runs: [Double] = []
-        for _ in 0 ..< 5 {
-            let caches = qwen.makeLayerCaches()
-            let cmd = Device.shared.makeCommandBuffer()
-            _ = qwen.forwardMany(
-                tokenIds: encoded, startPosition: 0,
-                caches: caches, on: cmd, device: Device.shared)
-            cmd.commit()
-            await cmd.completed()
+        // 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 t0 = Date()
-            for j in 0 ..< nSteps {
-                _ = qwen.forward(
-                    tokenId: 0, position: 1024 + j, caches: caches)
-            }
-            runs.append(Date().timeIntervalSince(t0))
+        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.completed()
+        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)
         }
-        runs.sort()
-        let median = runs[runs.count / 2]
-        let tps = Double(nSteps) / median
+        let decodeS = Date().timeIntervalSince(decodeT0)
+        let decodeTps = Double(nDecode) / decodeS
         print(
-            "Qwen3.5-35B-A3B decode T=1 after T=1024 prefill: "
-                + "runs=\(runs.map { String(format: "%.3f", $0) })s "
-                + "median=\(String(format: "%.3f", median))s → "
-                + "\(String(format: "%.2f", tps)) tps")
+            "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 {

From 09830664849eff4835d9e29a4c8eb117f952e342 Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 16:30:53 -0500
Subject: [PATCH 46/54] fix(tests): portable cmd.awaitCompletion() bridges
 Xcode 16.4 + 26.5
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

CI's macos-latest is on Xcode 16.4 which lacks the async `MTLCommand-
Buffer.completed()` extension Apple added in Xcode 26.x. Local dev box
on Xcode 26.5 ships `completed()` AND flags the long-standing sync
`waitUntilCompleted()` as "unavailable from asynchronous contexts" in
async test functions. So neither single API works on both toolchains.

Adds `MTLCommandBuffer.awaitCompletion()` to `Tests/Helpers/CommonTest-
Helpers.swift` — bridges `addCompletedHandler` through `withChecked-
Continuation`. Works on every Metal SDK from macOS 10.11 onward, no
compiler complaints on either toolchain, no real-thread blocking.

Replaces 11 call sites in:
- Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift
- Tests/ModelIntegrationTests/Text/Qwen36TextIntegrationTests.swift
- Tests/ModelIntegrationTests/Text/SpecDecodeVerifyCostBench.swift

Pre-existing CI break that surfaced after the metaltile dependency PRs
landed in #193-#201 unblocked compilation past the previous
"MetalTileKernels has no member ..." cascade.
---
 Tests/Helpers/CommonTestHelpers.swift           | 17 +++++++++++++++++
 .../Text/Qwen35MoEBenchIntegrationTests.swift   |  8 ++++----
 .../Text/Qwen36TextIntegrationTests.swift       |  6 +++---
 .../Text/SpecDecodeVerifyCostBench.swift        |  8 ++++----
 4 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/Tests/Helpers/CommonTestHelpers.swift b/Tests/Helpers/CommonTestHelpers.swift
index e3daaf41..fe252742 100644
--- a/Tests/Helpers/CommonTestHelpers.swift
+++ b/Tests/Helpers/CommonTestHelpers.swift
@@ -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<Void, Never>) in
+            self.addCompletedHandler { _ in cont.resume() }
+        }
+    }
+}
 
 // MARK: - ModelLoadLock
 
diff --git a/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift b/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift
index 3c861f8f..c60bd0c1 100644
--- a/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift
+++ b/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift
@@ -103,7 +103,7 @@ struct Qwen35MoEBenchIntegrationTests {
             tokenIds: encoded, startPosition: 0,
             caches: caches, on: cmd, device: Device.shared)
         cmd.commit()
-        await cmd.completed()
+        await cmd.awaitCompletion()
         let dt = Date().timeIntervalSince(t0)
         print("Qwen3.5-35B-A3B forwardMany T=8: \(String(format: "%.3f", dt))s")
     }
@@ -181,7 +181,7 @@ struct Qwen35MoEBenchIntegrationTests {
             tokenIds: encoded, startPosition: 0,
             caches: caches, on: cmd, device: Device.shared)
         cmd.commit()
-        await cmd.completed()
+        await cmd.awaitCompletion()
         let prefillS = Date().timeIntervalSince(prefillT0)
         let prefillTps = Double(T) / prefillS
         print(
@@ -242,7 +242,7 @@ struct Qwen35MoEBenchIntegrationTests {
                 tokenIds: encoded, startPosition: 0,
                 caches: warmCachesB, on: warmCmd, device: Device.shared)
             warmCmd.commit()
-            await warmCmd.completed()
+            await warmCmd.awaitCompletion()
         }
 
         // Per-token loop baseline (5 runs, median).
@@ -268,7 +268,7 @@ struct Qwen35MoEBenchIntegrationTests {
                 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()
diff --git a/Tests/ModelIntegrationTests/Text/Qwen36TextIntegrationTests.swift b/Tests/ModelIntegrationTests/Text/Qwen36TextIntegrationTests.swift
index 0f67c276..f077ffee 100644
--- a/Tests/ModelIntegrationTests/Text/Qwen36TextIntegrationTests.swift
+++ b/Tests/ModelIntegrationTests/Text/Qwen36TextIntegrationTests.swift
@@ -220,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
         }
 
@@ -250,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()
@@ -297,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/SpecDecodeVerifyCostBench.swift b/Tests/ModelIntegrationTests/Text/SpecDecodeVerifyCostBench.swift
index 9fb95334..08406ef6 100644
--- a/Tests/ModelIntegrationTests/Text/SpecDecodeVerifyCostBench.swift
+++ b/Tests/ModelIntegrationTests/Text/SpecDecodeVerifyCostBench.swift
@@ -90,7 +90,7 @@ struct SpecDecodeVerifyCostBench {
             tokenIds: inputIds, startPosition: promptLen,
             caches: caches, on: warmCmd, device: device)
         warmCmd.commit()
-        await warmCmd.completed()
+        await warmCmd.awaitCompletion()
         caches.restoreAll(from: snap0, device: device)
 
         // 16 timed iters of each path, restoring caches between each.
@@ -115,7 +115,7 @@ struct SpecDecodeVerifyCostBench {
                 tokenIds: inputIds, startPosition: promptLen,
                 caches: caches, on: cmd, device: device)
             cmd.commit()
-            await cmd.completed()
+            await cmd.awaitCompletion()
             batchedTimes.append(Date().timeIntervalSince(t0))
             caches.restoreAll(from: snap0, device: device)
         }
@@ -154,7 +154,7 @@ struct SpecDecodeVerifyCostBench {
             tokenIds: inputIds2, startPosition: promptLen,
             caches: caches, on: warmCmd2, device: device)
         warmCmd2.commit()
-        await warmCmd2.completed()
+        await warmCmd2.awaitCompletion()
         caches.restoreAll(from: snap2, device: device)
 
         var singleT2: [Double] = []
@@ -175,7 +175,7 @@ struct SpecDecodeVerifyCostBench {
                 tokenIds: inputIds2, startPosition: promptLen,
                 caches: caches, on: cmd, device: device)
             cmd.commit()
-            await cmd.completed()
+            await cmd.awaitCompletion()
             batchedT2.append(Date().timeIntervalSince(t0))
             caches.restoreAll(from: snap2, device: device)
         }

From 2a03114009e3703c27eaf82d27ebb170f960571d Mon Sep 17 00:00:00 2001
From: TheTom <tturney1@gmail.com>
Date: Tue, 26 May 2026 21:50:07 -0500
Subject: [PATCH 47/54] test: gate bm64_mpp bf16 cells on Apple GPU Family 9+

The mt_moe_gather_qmm_mma_int4_bm64_mpp_bf16 kernel uses MPP cooperative
tensor intrinsics (simdgroup_matrix on cooperative-tensor types) that
require Apple GPU Family 9 or later. On Family 7 (M1, CI runner) the
kernel compiles but emits zeroed output, producing cos=0.0 vs the bm16
and m1 references. The kernel is independently covered by upstream
metaltile's GPU correctness suite on adequate hardware.

Gates all 7 bf16 cells in MoEBgemmBm64MppTests with .enabled(if:) using
a static MTLDevice.supportsFamily(.apple9) probe. M5 Max (Family 9+)
runs untouched; M1 runners skip with a clear reason string.
---
 .../Benchmark/MoEBgemmBm64MppTests.swift      | 40 +++++++++++++++----
 1 file changed, 33 insertions(+), 7 deletions(-)

diff --git a/Tests/FFAITests/Benchmark/MoEBgemmBm64MppTests.swift b/Tests/FFAITests/Benchmark/MoEBgemmBm64MppTests.swift
index 946aec65..8e254826 100644
--- a/Tests/FFAITests/Benchmark/MoEBgemmBm64MppTests.swift
+++ b/Tests/FFAITests/Benchmark/MoEBgemmBm64MppTests.swift
@@ -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

From 4d7d4de3fdca1b01285ec65a95697d6d71e75c76 Mon Sep 17 00:00:00 2001
From: Eric Kryski <599019+ekryski@users.noreply.github.com>
Date: Tue, 26 May 2026 23:34:38 -0600
Subject: [PATCH 48/54] chore(qwen35): remove dead let _ = and unused dtype
 bindings

- Drop `let _ = resultFlat` immediately before `return resultFlat` in
  Qwen35DecoderLayer.decodeMany; it was a stale leftover.
- Drop unused `let dt = hFlat.dtype` + `let dtBytes = dt.byteSize` pair
  in both Qwen35GDNLayer.decodeMany and Qwen35AttentionLayer.decodeMany.
---
 Sources/FFAI/Models/Text/Qwen3xText.swift | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/Sources/FFAI/Models/Text/Qwen3xText.swift b/Sources/FFAI/Models/Text/Qwen3xText.swift
index 68f65b74..163e50bb 100644
--- a/Sources/FFAI/Models/Text/Qwen3xText.swift
+++ b/Sources/FFAI/Models/Text/Qwen3xText.swift
@@ -2492,9 +2492,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,
@@ -2646,9 +2643,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,
@@ -2771,7 +2765,6 @@ private func qwen35ApplyFFNMany(
         let addCmd = device.makeCommandBuffer()
         let resultFlat = Ops.add(postMix, ffnOut, on: addCmd)
         addCmd.commit()
-        let _ = resultFlat
         return resultFlat
     }
 }

From 90aeafcc37974145ead09f4aa9c9cd2e9b69cbc3 Mon Sep 17 00:00:00 2001
From: Eric Kryski <599019+ekryski@users.noreply.github.com>
Date: Tue, 26 May 2026 23:36:08 -0600
Subject: [PATCH 49/54] fix(ops): cross-check K/V packed widths in batched*
 fast wrappers

The batchedQkvQ*Fast and batched4*Fast wrappers derive `in_dim` from
`wQ.shape[1] * 8` (or `wA.shape[1] * 8` for the 4-way) and pass that
single value to the kernel. The kernel walks every weight tensor at
that packed width, so a mis-packed wK/wV (or wB/wC/wD) would silently
miscompute its branch's output rather than fail at dispatch time.

Add an explicit precondition that all packed widths match the leading
weight tensor, with an error message naming the mismatched widths so
the failure mode is obvious if it ever fires.

Affects:
- Ops.batchedQkvQgemvInt4Fast
- Ops.batchedQkvQmmFast
- Ops.batched4QgemvInt4Fast
- Ops.batched4QmmFast
---
 Sources/FFAI/Ops/Ops.swift | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift
index 0e6b2a91..02e887c5 100644
--- a/Sources/FFAI/Ops/Ops.swift
+++ b/Sources/FFAI/Ops/Ops.swift
@@ -1896,6 +1896,13 @@ public enum Ops {
             "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)")
@@ -2003,6 +2010,13 @@ public enum Ops {
             "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)")
@@ -2122,6 +2136,14 @@ public enum Ops {
             "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)")
@@ -2244,6 +2266,14 @@ public enum Ops {
             "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)")

From ac3ce85e457c2a148de76e67c1fdb240b81c0be5 Mon Sep 17 00:00:00 2001
From: Eric Kryski <599019+ekryski@users.noreply.github.com>
Date: Tue, 26 May 2026 23:37:35 -0600
Subject: [PATCH 50/54] docs(threading): document single-inference-per-instance
 contract
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

MoELayer and Qwen35{AttentionMixer,GDNMixer,MoEFFN} each hold a growing
set of per-instance scratch tensors that are mutated by `decode` /
`forward` / `forwardMany` without internal synchronisation. The
implicit contract — one inference per layer instance at a time — is
satisfied today by the per-token serial model loop, but is not
expressed in the API.

Make it explicit on every affected class header. Also tighten the
contract note on `MoELayer.decodeGPURouter` to call out that the
returned tensor aliases a scratch the GPU is still writing to via the
just-committed cmd, so a re-entrant call would race the in-flight
work.
---
 Sources/FFAI/Models/MoELayer.swift        | 18 ++++++++++++++++++
 Sources/FFAI/Models/Text/Qwen3xText.swift | 20 ++++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/Sources/FFAI/Models/MoELayer.swift b/Sources/FFAI/Models/MoELayer.swift
index 36a111d4..6ae6de66 100644
--- a/Sources/FFAI/Models/MoELayer.swift
+++ b/Sources/FFAI/Models/MoELayer.swift
@@ -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
@@ -1313,6 +1322,15 @@ public final class MoELayer: Module, DecoderLayer {
         // (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
     }
diff --git a/Sources/FFAI/Models/Text/Qwen3xText.swift b/Sources/FFAI/Models/Text/Qwen3xText.swift
index 163e50bb..caf54c31 100644
--- a/Sources/FFAI/Models/Text/Qwen3xText.swift
+++ b/Sources/FFAI/Models/Text/Qwen3xText.swift
@@ -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
@@ -1130,6 +1137,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
@@ -1866,6 +1879,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

From 272d8f122a5eae3632121ae010e7f3a24876cdf6 Mon Sep 17 00:00:00 2001
From: Eric Kryski <599019+ekryski@users.noreply.github.com>
Date: Tue, 26 May 2026 23:38:10 -0600
Subject: [PATCH 51/54] docs(kvcache): document cachedSnapshot reuse contract
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

GDNStateCache.snapshot() and ConvStateCache.snapshot() both return a
single per-instance scratch tensor that is overwritten on subsequent
calls. The current spec-decode driver does snapshot → verify → restore
once per layer per verify step so this is safe in practice, but the
contract is invisible to a caller reading just the method signature.

Spell it out: a second snapshot() before the matching restore() will
trample the prior contents. Future nested or concurrent snapshot usage
would need a different design (per-call allocation, or a small ring).
---
 Sources/FFAI/KVCache/ConvStateCache.swift | 5 +++++
 Sources/FFAI/KVCache/GDNStateCache.swift  | 5 +++++
 2 files changed, 10 insertions(+)

diff --git a/Sources/FFAI/KVCache/ConvStateCache.swift b/Sources/FFAI/KVCache/ConvStateCache.swift
index 68e56104..f769b8a1 100644
--- a/Sources/FFAI/KVCache/ConvStateCache.swift
+++ b/Sources/FFAI/KVCache/ConvStateCache.swift
@@ -69,6 +69,11 @@ public final class ConvStateCache: @unchecked Sendable {
 
     /// 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
diff --git a/Sources/FFAI/KVCache/GDNStateCache.swift b/Sources/FFAI/KVCache/GDNStateCache.swift
index 59150f7b..2fb65ec3 100644
--- a/Sources/FFAI/KVCache/GDNStateCache.swift
+++ b/Sources/FFAI/KVCache/GDNStateCache.swift
@@ -142,6 +142,11 @@ public final class GDNStateCache: LayerCacheProtocol, @unchecked Sendable {
     /// 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 {

From 8f3c409de7acf2e3d6732a0b728ce7baf25d20b2 Mon Sep 17 00:00:00 2001
From: Eric Kryski <599019+ekryski@users.noreply.github.com>
Date: Tue, 26 May 2026 23:44:59 -0600
Subject: [PATCH 52/54] docs: strip internal iteration markers from code
 comments
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Remove all "ITER N" / "(PR10)" / "PR10" markers from comments in
the model and op-wiring code. These were references to an authoring
agent's internal iteration log and don't help a future reader — they
made comments harder to skim and left a trail of stale cross-refs.

Where the surrounding comment depended on the marker for grammar or
context, tighten the prose so the remaining sentence still reads
cleanly. No behavioural changes; the comments still describe the
same fast-paths, gating env vars, and dispatch shapes.
---
 Sources/FFAI/Models/MoELayer.swift            |  70 +++--
 Sources/FFAI/Models/Text/Qwen3xText.swift     | 250 +++++++++---------
 Sources/MetalTileSwift/MetalTileLibrary.swift |   6 +-
 3 files changed, 160 insertions(+), 166 deletions(-)

diff --git a/Sources/FFAI/Models/MoELayer.swift b/Sources/FFAI/Models/MoELayer.swift
index 6ae6de66..90425d0f 100644
--- a/Sources/FFAI/Models/MoELayer.swift
+++ b/Sources/FFAI/Models/MoELayer.swift
@@ -399,32 +399,32 @@ public final class MoELayer: Module, DecoderLayer {
     public let enableBGEMM: Bool
     public let useBm8Env: Bool
     public let useM1Env: Bool
-    /// ITER 64 (PR10): additional env-flag caches.
+    /// Additional env-flag caches.
     let noDequantGate: Bool
     let gpuRouterEnabled: Bool
     let forceBGEMMAtT1: Bool
     let noBGEMMBm64: Bool
 
-    /// ITER 15/42 (PR10): lazy dequant of the gate weight + cached output.
-    /// Skipped when `gate.inner` is not a 8-bit QuantizedLinear or env
+    /// 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?
 
-    /// ITER 41 (PR10): cached scratches for the per-expert weighted-sum
-    /// loop at decode T=1.
+    /// Cached scratches for the per-expert weighted-sum loop at
+    /// decode T=1.
     private var accumulatorScratch: Tensor?
     private var topKScalarsBuf: MTLBuffer?
-    /// ITER 32 / 36 / 38 (PR10): cached per-expert g/u/inner/down outputs.
+    /// Cached per-expert g/u/inner/down outputs.
     private var expertGScratches: [Tensor] = []
     private var expertUScratches: [Tensor] = []
     private var expertInnerScratches: [Tensor] = []
     private var expertOutScratches: [Tensor] = []
-    /// ITER 56 (PR10): 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.
+    /// 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?
 
@@ -488,9 +488,9 @@ public final class MoELayer: Module, DecoderLayer {
         self.useBm8Env = env["FFAI_MOE_BGEMM_BM8"] != nil
         self.useM1Env = env["FFAI_MOE_M1"] != nil
         self.noDequantGate = env["FFAI_MOE_NO_DEQUANT_GATE"] != nil
-        // ITER 80 (PR10): default GPU router ON now that ITER 72's
-        // end-to-end correctness pinning is in place across bf16/f16/f32.
-        // Opt out with `FFAI_MOE_GPU_ROUTER=0`. The win: 40 ×
+        // 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
@@ -539,9 +539,9 @@ public final class MoELayer: Module, DecoderLayer {
         // ── 1. Gate gemv on the caller's command buffer ──────────────
         let logitsTensor = gate(h, on: cmd)
 
-        // ── 2a. ITER 56 (PR10): GPU MoE router fast path ─────────────
-        // When `FFAI_MOE_GPU_ROUTER=1` (default ON per ITER 80) AND the
-        // layer has stacked int4 expert layout AND routing matches
+        // ── 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-
@@ -691,9 +691,9 @@ public final class MoELayer: Module, DecoderLayer {
         // gateLogitsAll shape: [T, nExperts]
 
         // ── 2. Routing — GPU fast path or CPU sync ──────────────────────
-        // ITER 92 (PR10): 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
+        // 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 +
@@ -858,9 +858,9 @@ public final class MoELayer: Module, DecoderLayer {
         //     2.69× T=32 win.
         //   - mTotal ≤ 8 + `FFAI_MOE_BGEMM_BM8=1` → bm8 (decode T=1
         //     fallback).
-        // ITER 73 (PR10): bm64_mpp default-on for mTotal ≥ 1024 (refreshed
-        // bench: crossover sits between mTotal=512 bm16 +19% and
-        // mTotal=1024 bm64 +18%). Opt out via FFAI_MOE_BGEMM_NO_BM64=1.
+        // 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:
@@ -1147,13 +1147,12 @@ public final class MoELayer: Module, DecoderLayer {
         return downProj(inner, on: cmd)
     }
 
-    /// ITER 56 (PR10): GPU MoE router decode path. Replaces the
+    /// 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,
-    /// ITER 94).
+    /// `moeDownSwigluAccumInt4Chain8` kernel when constraints allow).
     ///
     /// Predicated by `decode` callers — the layer must have stackedInt4
     /// experts with matching dtype, router must be softmax-then-topK
@@ -1202,10 +1201,10 @@ public final class MoELayer: Module, DecoderLayer {
             normTopkProb: router.normTopKProb,
             on: cmd)
 
-        // ITER 58 (PR10): 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.
+        // 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 {
@@ -1248,14 +1247,13 @@ public final class MoELayer: Module, DecoderLayer {
             inputs: ins, expertIndices: idxs, outputs: outs0,
             groupSize: stacked.groupSize, on: cmd)
 
-        // ITER 94 (PR10): 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 (ITER 29/30/31, 38/39, ITER 47) for any larger configs.
+        // 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!
diff --git a/Sources/FFAI/Models/Text/Qwen3xText.swift b/Sources/FFAI/Models/Text/Qwen3xText.swift
index caf54c31..d32ba889 100644
--- a/Sources/FFAI/Models/Text/Qwen3xText.swift
+++ b/Sources/FFAI/Models/Text/Qwen3xText.swift
@@ -777,7 +777,7 @@ public final class Qwen35DenseMLP: Module {
         return out
     }
 
-    /// down(silu(gate(x)) * up(x)). ITER 16: use fused Ops.swiglu —
+    /// 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)
@@ -824,14 +824,14 @@ public final class Qwen35MoEFFN: Module {
     let sharedExpertGate: AnyLinear
     let hidden: Int
 
-    /// ITER 26/82 (PR10): 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.
+    /// 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?
-    /// ITER 34: cached `result` output for the GPU fan-out.
+    /// Cached `result` output for the GPU fan-out.
     private var sharedResultScratch: Tensor?
 
     init(
@@ -871,8 +871,8 @@ public final class Qwen35MoEFFN: Module {
     /// shared expert + the final add run on fresh private buffers, so
     /// the returned tensor never depends on the now-dead `cmd`.
     ///
-    /// ITER 66 (PR10): optional `residual` parameter folds the post-FFN
-    /// residual add into the final `sigmoidScalarFMA` dispatch (becomes
+    /// 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)`.
@@ -888,12 +888,10 @@ public final class Qwen35MoEFFN: Module {
             cmd: cmd, device: device)
         // Shared expert on a fresh buffer.
         let work = device.makeCommandBuffer()
-        // ITER 26: batch sharedGate + sharedUp qmms onto one shared
-        // encoder via dequantGemvInt4Two when both are int4.
-        // ITER 82: also fold sharedExpertGate (an int4 qmm producing a
-        // single scalar from the same `xNorm`) into the SAME shared
-        // encoder via dequantGemvInt4Three. Saves another encoder
-        // begin/end per MoE layer × 40 ≈ 40 pairs / decode token.
+        // 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
@@ -958,13 +956,13 @@ public final class Qwen35MoEFFN: Module {
             su = sharedUpProj(xNorm, on: work)
             gateLogit = sharedExpertGate(xNorm, on: work)
         }
-        // ITER 16: fused SwiGLU.
+        // Fused SwiGLU.
         let sharedInner = Ops.swiglu(gate: sg, up: su, on: work)
         let sharedOut = sharedDownProj(sharedInner, on: work)
         work.commit()
 
         // GPU fan-out: `out = routed + sigmoid(gateLogit) * sharedOut`
-        // (+ residual when ITER 66 is in play) in one dispatch.
+        // (+ residual when the caller passes one) in one dispatch.
         let fmaCmd = device.makeCommandBuffer()
         if sharedResultScratch == nil
             || sharedResultScratch!.dtype != routed.dtype
@@ -975,7 +973,7 @@ public final class Qwen35MoEFFN: Module {
         }
         let result = sharedResultScratch!
         if let residual = residual {
-            // ITER 66: fused 4-input. result = residual + routed +
+            // Fused 4-input. result = residual + routed +
             // sigmoid(gateLogit) · sharedOut. Saves 1 dispatch + 1
             // [hidden] DRAM roundtrip vs sigmoidScalarFMA then Ops.add.
             Ops.sigmoidScalarFMAResidual(
@@ -1013,7 +1011,7 @@ 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)
-        // ITER 16: fused SwiGLU.
+        // Fused SwiGLU.
         let sharedInnerAll = Ops.swiglu(gate: sgAll, up: suAll, on: work)
         let sharedOutAll = sharedDownProj.callMany(
             sharedInnerAll, t: t,
@@ -1206,12 +1204,12 @@ public final class Qwen35GDNMixer: Module {
     let yF32Scratch: Tensor
     let yGatedScratch: Tensor
 
-    /// ITER 23 (PR10): 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
+    /// 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 via the
-    /// residency set in init (ITER 27).
+    /// 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,
@@ -1224,8 +1222,8 @@ public final class Qwen35GDNMixer: Module {
     private var zScratch: Tensor?
     private var bRawScratch: Tensor?
     private var aRawScratch: Tensor?
-    /// ITER 95 (PR10): single-dispatch `ffai_batched_4_qgemv_fast` fast
-    /// path. Constraints: in_dim%512==0, each out_dim%8==0, groupSize=64.
+    /// 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(
@@ -1275,9 +1273,9 @@ public final class Qwen35GDNMixer: Module {
         self.aLogTF32 = makeF32Tensor35(aLog, device: device)
         self.dtBiasTF32 = makeF32Tensor35(dtBias, device: device)
         self.epsBufFused = makeF32Tensor35([eps], device: device)
-        // ITER 19 (PR10): 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.
+        // 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
 
@@ -1292,9 +1290,10 @@ public final class Qwen35GDNMixer: Module {
             dtype: .f32, device: device)
         self.yGatedScratch = Tensor.empty(shape: [valueDim], dtype: dtype, device: device)
 
-        // ITER 23 (PR10): set up the batched 4-projection input fast
-        // path when all 4 inProj are QuantizedLinear int4 with same
-        // groupSize. ITER 27: pin scratches in the residency set.
+        // 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,
@@ -1311,7 +1310,7 @@ public final class Qwen35GDNMixer: Module {
                 aW: aQL.weight, aS: aQL.scales, aB: aQL.biases,
                 groupSize: qQL.groupSize
             )
-            // ITER 95: single-dispatch fast path eligibility.
+            // Single-dispatch fast path eligibility.
             let inDim = qQL.weight.shape[1] * 8
             if inDim % 512 == 0
                 && convDim % 8 == 0 && valueDim % 8 == 0
@@ -1329,7 +1328,7 @@ public final class Qwen35GDNMixer: Module {
                 shape: [numValueHeads], dtype: dtype, device: device)
             self.aRawScratch = Tensor.empty(
                 shape: [numValueHeads], dtype: dtype, device: device)
-            // ITER 27: pin scratch buffers + fused-path constants so Metal
+            // Pin scratch buffers + fused-path constants so Metal
             // skips per-encoder residency revalidation.
             device.markWeightsResident([
                 qkvScratch!.buffer, zScratch!.buffer,
@@ -1375,10 +1374,10 @@ public final class Qwen35GDNMixer: Module {
         cmd: MTLCommandBuffer, device: Device
     ) -> Tensor {
         // ── GPU phase 1: projections + conv + SiLU ────────────────────
-        // ITER 23 (PR10): 4-projection shared-encoder fast path when all
-        // inProj are int4 QuantizedLinear with same groupSize. Saves 3
-        // encoder begin/end pairs per GDN layer × 30 ≈ 1.5 ms/token.
-        // ITER 95: single-dispatch fast path via `batched4QgemvInt4Fast`
+        // 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
@@ -1417,11 +1416,11 @@ public final class Qwen35GDNMixer: Module {
             x: qkv, w: convW, b: convB,
             state: cache.conv.state, into: convOutScratch,
             nChannels: convDim, kernelSize: convKernel, on: cmd)
-        // ITER 81 (PR10): 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.
+        // 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
@@ -1457,10 +1456,10 @@ 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.
-            // ITER 81 (PR10): 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.
+            // 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,
@@ -1651,8 +1650,8 @@ public final class Qwen35GDNMixer: Module {
             "Qwen35GDNMixer.forwardMany: xNormFlat size \(xNormFlat.elementCount) ≠ T·hidden = \(t * hidden)"
         )
 
-        // PR10: chunked prep+recurrence is the default route. Replaces
-        // the per-token T-loop body with one `Ops.gatedDeltaPrepChunk`
+        // 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 {
@@ -1669,9 +1668,8 @@ 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)
 
-        // ITER 87 (PR10) preview: a_raw / b_raw → f32 once via
-        // shared-encoder `castToF32Two`, then per-row slices alias into
-        // the pre-cast f32 storage.
+        // 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(
@@ -1749,7 +1747,7 @@ public final class Qwen35GDNMixer: Module {
     /// 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 ITER 99's batched
+    /// 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.
@@ -1773,8 +1771,8 @@ 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)
 
-        // ITER 87 (PR10): a_raw / b_raw → f32 once on ONE shared encoder
-        // via `castToF32Two`. Saves one encoder begin/end pair per GDN
+        // 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)
@@ -1786,7 +1784,7 @@ public final class Qwen35GDNMixer: Module {
             on: cmd)
 
         // ── Conv1d + silu+cast → `[T, convDim]` f32 staging buffer ──────
-        // ITER 99 (PR10): when conv_kernel = 4, the batched
+        // 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
@@ -1856,8 +1854,8 @@ public final class Qwen35GDNMixer: Module {
         for _ in 0 ..< t { cache.advance() }
 
         // ── Mixer norm T-batched (one dispatch across T·Hv rows) ────────
-        // ITER (PR10): `gatedMixerNormMany` collapses T per-row
-        // dispatches into one shared-encoder dispatch.
+        // `gatedMixerNormMany` collapses T per-row dispatches into one
+        // shared-encoder dispatch.
         let yGatedAll = Tensor.empty(
             shape: [t * valueDim], dtype: dt, device: device)
         Ops.gatedMixerNormMany(
@@ -1894,14 +1892,14 @@ public final class Qwen35AttentionMixer: Module {
     let ropeTheta: Float
     let attnOutputGate: Bool
     let scale: Float
-    /// ITER 22 (PR10): 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.
+    /// 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, ...
-    /// ITER 24 (PR10): batched QKV-projection fast-path detection. When
-    /// all 3 projections (q, k, v) are QuantizedLinear int4 with same
+    /// 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:
@@ -1914,25 +1912,25 @@ public final class Qwen35AttentionMixer: Module {
     private var qOutScratch: Tensor?
     private var kOutScratch: Tensor?
     private var vOutScratch: Tensor?
-    /// ITER 78 (PR10): fused-QKV single-dispatch fast path. When set,
-    /// `forward` uses `Ops.batchedQkvQgemvInt4Fast` (1 dispatch) instead
-    /// of `dequantGemvInt4Three` (3 dispatches on shared encoder).
+    /// 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?
-    /// ITER 35 (PR10): cached q_norm + k_norm outputs (was fresh per call).
+    /// Cached q_norm + k_norm outputs (was fresh per call).
     private var qNormedScratch: Tensor?
     private var kNormedScratch: Tensor?
-    /// ITER 40 (PR10): cached sliceHeadHalves35 gather outputs.
+    /// Cached sliceHeadHalves35 gather outputs.
     private var queriesScratch: Tensor?
     private var gateSliceScratch: Tensor?
-    /// ITER 43 (PR10): cached sdpaDecode output + sigmoidMul output.
+    /// Cached sdpaDecode output + sigmoidMul output.
     private var attnOutScratch: Tensor?
     private var attnFlatGatedScratch: Tensor?
-    /// ITER 53 (PR10): 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.
+    /// 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?
@@ -1967,7 +1965,7 @@ public final class Qwen35AttentionMixer: Module {
         self.attnOutputGate = attnOutputGate
         self.scale = 1.0 / Float(Double(headDim).squareRoot())
         if attnOutputGate {
-            // ITER 22: pre-build the row-index tensors for the gate split
+            // 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)
@@ -1987,13 +1985,13 @@ public final class Qwen35AttentionMixer: Module {
                 buffer: firstBuf, offset: 0, shape: [nHeads], dtype: .u32)
             self.sliceSecondIdx = Tensor(
                 buffer: secondBuf, offset: 0, shape: [nHeads], dtype: .u32)
-            // ITER 28: pin idx scratches to the residency set.
+            // Pin idx scratches to the residency set.
             device.markWeightsResident([firstBuf, secondBuf])
         } else {
             self.sliceFirstIdx = nil
             self.sliceSecondIdx = nil
         }
-        // ITER 24: detect batched-qkv fast path.
+        // Detect batched-qkv fast path.
         if let qQL = qProj.inner as? QuantizedLinear,
             let kQL = kProj.inner as? QuantizedLinear,
             let vQL = vProj.inner as? QuantizedLinear,
@@ -2006,8 +2004,8 @@ public final class Qwen35AttentionMixer: Module {
                 vW: vQL.weight, vS: vQL.scales, vB: vQL.biases,
                 groupSize: qQL.groupSize
             )
-            // ITER 78: fused single-dispatch QKV gemv. Constraints:
-            // in_dim multiple of 512, each out_dim multiple of 8,
+            // 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]
@@ -2044,9 +2042,9 @@ 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.
-        // ITER 24 (PR10): batch q+k+v projections into one encoder when
-        // all 3 are QuantizedLinear int4. ITER 78: when constraints
-        // match, fuse all three into a single dispatch.
+        // 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
@@ -2055,7 +2053,7 @@ public final class Qwen35AttentionMixer: Module {
             let kDim = bq.kW.shape[0]
             let vDim = bq.vW.shape[0]
             if fusedQKVEligible {
-                // ITER 78: ONE dispatch into a [qDim+kDim+vDim] scratch.
+                // 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
@@ -2080,7 +2078,7 @@ public final class Qwen35AttentionMixer: Module {
                     outQ: qDim, outK: kDim, outV: vDim,
                     on: cmd, into: qkvFusedScratch!)
             } else {
-                // Legacy 3-dispatch shared-encoder path (ITER 24).
+                // Legacy 3-dispatch shared-encoder path.
                 if qOutScratch == nil || qOutScratch!.elementCount != qDim {
                     qOutScratch = Tensor.empty(
                         shape: [qDim], dtype: xNorm.dtype, device: device)
@@ -2113,7 +2111,7 @@ public final class Qwen35AttentionMixer: Module {
         if attnOutputGate {
             // Layout is [nHeads, 2 · headDim] — per head the first
             // `headDim` is the query, the next `headDim` the gate.
-            // ITER 22: use cached idx tensors.
+            // Use cached idx tensors.
             let q2 = qOut.reshaped(to: [nHeads, 2 * headDim])
             let table = q2.reshaped(to: [nHeads * 2, headDim])
             if queriesScratch == nil
@@ -2125,8 +2123,8 @@ public final class Qwen35AttentionMixer: Module {
                 gateSliceScratch = Tensor.empty(
                     shape: [nHeads, headDim], dtype: qOut.dtype, device: device)
             }
-            // ITER 65: batched two-gather on shared encoder, saves
-            // 1 encoder begin/end pair per attn layer × 10 = ~170 µs/token.
+            // 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!,
@@ -2142,7 +2140,7 @@ public final class Qwen35AttentionMixer: Module {
         let v = vPre
 
         // Per-head q_norm / k_norm (weighted RMSNorm over headDim).
-        // ITER 12 + ITER 35: shared encoder + cached scratches.
+        // Shared encoder + cached scratches.
         if qNormedScratch == nil
             || qNormedScratch!.elementCount != queries.elementCount
             || qNormedScratch!.dtype != queries.dtype
@@ -2162,7 +2160,7 @@ public final class Qwen35AttentionMixer: Module {
             on: cmd)
 
         // Partial RoPE — rotate only the first `rotaryDim` dims of each
-        // head, in place. ITER 21: q + k on shared encoder.
+        // head, in place. q + k on shared encoder.
         Ops.ropePartialTwo(
             qNormed, kNormed, position: position,
             headDim: headDim, rotaryDim: rotaryDim,
@@ -2174,7 +2172,7 @@ public final class Qwen35AttentionMixer: Module {
             vFlat: v.reshaped(to: [nKVHeads, headDim]), on: cmd)
 
         let (cacheK, cacheV) = kv.prepareForAttention(on: cmd)
-        // ITER 43: cached sdpaDecode output + sigmoidMul output scratches.
+        // Cached sdpaDecode output + sigmoidMul output scratches.
         if attnOutScratch == nil
             || attnOutScratch!.elementCount != nHeads * headDim
             || attnOutScratch!.dtype != qNormed.dtype
@@ -2184,9 +2182,9 @@ public final class Qwen35AttentionMixer: Module {
             attnFlatGatedScratch = Tensor.empty(
                 shape: [nHeads * headDim], dtype: qNormed.dtype, device: device)
         }
-        // ITER 53: 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
+        // 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 {
@@ -2222,8 +2220,8 @@ public final class Qwen35AttentionMixer: Module {
         }
 
         // Gated output: attnOut * sigmoid(gate), then o_proj.
-        // ITER 11 + 40: fused sigmoidMul via mt_sigmoid_mul into cached
-        // scratch — saves 1 encoder per attn layer × 10 ≈ 170 µs/token.
+        // 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.sigmoidMul(
@@ -2264,12 +2262,12 @@ public final class Qwen35AttentionMixer: Module {
         let kvDim = nKVHeads * headDim
 
         // ── Projections — fused QKV qmm when constraints match ──────────
-        // ITER 88 (PR10): `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
+        // `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 (ITER 78).
+        // single-token path uses.
         let qOut: Tensor
         let kOut: Tensor
         let vOut: Tensor
@@ -2299,7 +2297,7 @@ public final class Qwen35AttentionMixer: Module {
         }
 
         // ── Gate split — one shared-encoder gather (vs T) when
-        // attnOutputGate. ITER 74: both halves on a shared encoder via
+        // attnOutputGate. Both halves run on a shared encoder via
         // sliceHeadHalvesManyBoth35.
         let queriesFlat: Tensor  // `[T, nHeads, headDim]` flat
         let gateFlat: Tensor?  // `[T, nHeads, headDim]` flat
@@ -2316,10 +2314,10 @@ public final class Qwen35AttentionMixer: Module {
         }
 
         // ── Q/K norm — one shared encoder over (T·nHeads) + (T·nKVHeads)
-        // rows. ITER 83 (PR10): collapse the two separate `rmsNormRows`
-        // dispatches into one `rmsNormRowsTwo` shared encoder, mirroring
-        // the single-token `forward` path (ITER 12+35). Saves 1 encoder
-        // begin/end per attn-layer prefill call.
+        // 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(
@@ -2332,10 +2330,8 @@ public final class Qwen35AttentionMixer: Module {
             on: cmd)
 
         // ── Batched RoPE over all T tokens on ONE shared encoder.
-        // ITER 84: pair Q+K RoPE for each row on the SAME shared encoder
-        // via `ropePartialTwo`. ITER 97: replace the T-loop of
-        // `Ops.ropePartialTwo` calls with ONE `Ops.ropePartialManyTwo`
-        // dispatch over all T tokens. Saves T-1 encoder begin/end pairs
+        // `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(
@@ -2350,7 +2346,7 @@ public final class Qwen35AttentionMixer: Module {
             headDim: headDim, rotaryDim: rotaryDim,
             thetaBase: ropeTheta, on: cmd)
 
-        // ── KV append — ITER 98: batched K+V cache append. Reserves T
+        // ── 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).
@@ -2361,8 +2357,8 @@ public final class Qwen35AttentionMixer: Module {
             kFlat: kNormed, vFlat: vOut, t: t,
             positions: appendPositions, on: cmd)
 
-        // ── SDPA — ITER 90 / 93: fuse the T-token causal attention into
-        // ONE `Ops.sdpaMulti` dispatch. d=128 wraps `ffai_sdpa_multi`;
+        // ── 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
@@ -2400,8 +2396,8 @@ public final class Qwen35AttentionMixer: Module {
             attnAll = attnAllScratch
         }
 
-        // ── Gated output * o_proj — ITER 14: use fused sigmoidMul on
-        // the prefill path (same kernel as the single-token ITER 11).
+        // ── 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.sigmoidMul(attnFlat, gateFlat, on: cmd)
@@ -2705,7 +2701,7 @@ public final class Qwen35AttentionLayer: Module, DecoderLayer {
     }
 }
 
-// ─── ITER 76 (PR10): finalNorm + lmHead fused dispatch helper ───────
+// ─── finalNorm + lmHead fused dispatch helper ───────────────────────
 //
 // Replaces the 2-dispatch chain `finalNorm(h)` + `lmHead(normed)` with
 // a single `Ops.rmsNormQgemvInt4Fast` dispatch when the lmHead's
@@ -2823,8 +2819,8 @@ private func qwen35ApplyFFN(
         }
         return result
     case .moe(let moe):
-        // Qwen35MoEFFN.forward commits `cmd`. ITER 66 (PR10): fold the
-        // post-FFN residual add into the final sigmoidScalarFMA via
+        // 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(
@@ -3003,8 +2999,8 @@ public final class Qwen35Model: LanguageModel {
         }
 
         // Final norm + lm_head queue onto the caller's pristine `cmd`.
-        // ITER 76 (PR10): fused rmsNorm+qgemv when lmHead is 4-bit
-        // QuantizedLinear with constraints; falls back otherwise.
+        // 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)
     }
@@ -3241,8 +3237,8 @@ public final class Qwen35Model: LanguageModel {
 
         // ── Final norm + lm_head ─────────────────────────────────────────
         // Two output shapes (see _forwardManyBatched docstring):
-        //   * returnAllLogits == false: logits of the LAST row only.
-        //     ITER 77 (PR10) fused via qwen35FinalNormLmHead.
+        //   * 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.
@@ -3310,8 +3306,8 @@ public final class Qwen35Model: LanguageModel {
             workCmd.waitUntilCompleted()
         }
 
-        // ITER 79 (PR10): fused finalNorm+lmHead on the VLM embed-input
-        // forward path too.
+        // Fused finalNorm+lmHead on the VLM embed-input forward path
+        // too.
         return qwen35FinalNormLmHead(
             h: h, finalNorm: finalNorm, lmHead: lmHead, on: cmd)
     }
@@ -3461,11 +3457,11 @@ private func sliceHeadHalves35(
     return gathered.reshaped(to: [nHeads * headDim])
 }
 
-/// T-batched both-halves variant (ITER 74). 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.
+/// 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,
diff --git a/Sources/MetalTileSwift/MetalTileLibrary.swift b/Sources/MetalTileSwift/MetalTileLibrary.swift
index 30184ebe..a8e1bb89 100644
--- a/Sources/MetalTileSwift/MetalTileLibrary.swift
+++ b/Sources/MetalTileSwift/MetalTileLibrary.swift
@@ -88,9 +88,9 @@ public final class MetalTileLibrary: @unchecked Sendable {
             return parsed
         }
         // 64 lets the Metal driver pipeline more cmd-buffers in flight
-        // before applying backpressure. PR10 ITER 20 measured the bump
-        // 16 → 64 as a 2-3% prefill T=512 win on Qwen3.6-A3B / M5 Max
-        // by absorbing the per-cmd-buffer encode latency across more
+        // 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

From d2367da557bd716e27dd46d836d6eaffcf5fda27 Mon Sep 17 00:00:00 2001
From: Eric Kryski <599019+ekryski@users.noreply.github.com>
Date: Tue, 26 May 2026 23:47:09 -0600
Subject: [PATCH 53/54] docs(copyright): credit Tom Turney on co-authored + new
 files
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

For files Tom co-authored (modifications on top of an existing Eric-
copyrighted file), add "and Tom Turney (@TheTom)" to the existing
copyright line.

For files Tom wrote from scratch, mark Tom as the sole copyright
holder. SpecDecode.swift, CacheSnapshot.swift, and DrafterTests.swift
were missing the standard Apache 2.0 boilerplate — add it now.

No code changes.
---
 Sources/FFAI/Device.swift                          |  2 +-
 Sources/FFAI/Generation/Drafter.swift              |  2 +-
 Sources/FFAI/Generation/SpecDecode.swift           | 14 ++++++++++++++
 Sources/FFAI/KVCache/CacheSnapshot.swift           | 14 ++++++++++++++
 Sources/FFAI/KVCache/ConvStateCache.swift          |  2 +-
 Sources/FFAI/KVCache/GDNStateCache.swift           |  2 +-
 Sources/FFAI/KVCache/KVCache.swift                 |  2 +-
 Sources/FFAI/Layers.swift                          |  2 +-
 Sources/FFAI/Loader/Model.swift                    |  2 +-
 Sources/FFAI/Models/Audio/VAD/FireRedVAD.swift     |  2 +-
 Sources/FFAI/Models/MoELayer.swift                 |  2 +-
 Sources/FFAI/Models/Text/Granite4Text.swift        |  2 +-
 Sources/FFAI/Models/Text/JambaText.swift           |  2 +-
 Sources/FFAI/Models/Text/LFM2Text.swift            |  2 +-
 Sources/FFAI/Models/Text/Qwen3xText.swift          |  2 +-
 Sources/FFAI/Ops/Ops.swift                         |  2 +-
 Sources/FFAI/Ops/OpsFused.swift                    |  2 +-
 Sources/FFAI/Ops/OpsValidation.swift               |  2 +-
 Sources/MetalTileSwift/MetalTileLibrary.swift      |  2 +-
 .../FFAITests/Benchmark/MoEBgemmBm64MppTests.swift |  2 +-
 Tests/FFAITests/Generation/DrafterTests.swift      | 14 ++++++++++++++
 Tests/FFAITests/Ops/OpsSpecialPathTests.swift      |  2 +-
 Tests/FFAITests/Ops/OpsTests.swift                 |  2 +-
 Tests/FFAITests/Ops/OpsValidationTests.swift       |  2 +-
 Tests/FFAITests/Ops/QuantizedOpsTests.swift        |  2 +-
 Tests/Helpers/CommonTestHelpers.swift              |  2 +-
 Tests/MetalTileSwiftTests/PSOCacheTests.swift      |  2 +-
 .../Text/Qwen35MoEBenchIntegrationTests.swift      |  2 +-
 .../Text/Qwen36TextIntegrationTests.swift          |  2 +-
 .../Text/SpecDecodeBenchTests.swift                |  2 +-
 .../Text/SpecDecodeVerifyCostBench.swift           |  2 +-
 31 files changed, 70 insertions(+), 28 deletions(-)

diff --git a/Sources/FFAI/Device.swift b/Sources/FFAI/Device.swift
index d6838c10..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.
diff --git a/Sources/FFAI/Generation/Drafter.swift b/Sources/FFAI/Generation/Drafter.swift
index a3165b1a..5443754c 100644
--- a/Sources/FFAI/Generation/Drafter.swift
+++ b/Sources/FFAI/Generation/Drafter.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.
diff --git a/Sources/FFAI/Generation/SpecDecode.swift b/Sources/FFAI/Generation/SpecDecode.swift
index 3c4f880a..be3d5a23 100644
--- a/Sources/FFAI/Generation/SpecDecode.swift
+++ b/Sources/FFAI/Generation/SpecDecode.swift
@@ -1,3 +1,17 @@
+// 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:
diff --git a/Sources/FFAI/KVCache/CacheSnapshot.swift b/Sources/FFAI/KVCache/CacheSnapshot.swift
index b16fc78d..d9c74011 100644
--- a/Sources/FFAI/KVCache/CacheSnapshot.swift
+++ b/Sources/FFAI/KVCache/CacheSnapshot.swift
@@ -1,3 +1,17 @@
+// 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
diff --git a/Sources/FFAI/KVCache/ConvStateCache.swift b/Sources/FFAI/KVCache/ConvStateCache.swift
index f769b8a1..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.
diff --git a/Sources/FFAI/KVCache/GDNStateCache.swift b/Sources/FFAI/KVCache/GDNStateCache.swift
index 2fb65ec3..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.
diff --git a/Sources/FFAI/KVCache/KVCache.swift b/Sources/FFAI/KVCache/KVCache.swift
index 505c93be..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.
diff --git a/Sources/FFAI/Layers.swift b/Sources/FFAI/Layers.swift
index 5ad8a750..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.
diff --git a/Sources/FFAI/Loader/Model.swift b/Sources/FFAI/Loader/Model.swift
index fb0eba39..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.
diff --git a/Sources/FFAI/Models/Audio/VAD/FireRedVAD.swift b/Sources/FFAI/Models/Audio/VAD/FireRedVAD.swift
index de623599..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.
diff --git a/Sources/FFAI/Models/MoELayer.swift b/Sources/FFAI/Models/MoELayer.swift
index 90425d0f..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.
diff --git a/Sources/FFAI/Models/Text/Granite4Text.swift b/Sources/FFAI/Models/Text/Granite4Text.swift
index adcb27e3..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.
diff --git a/Sources/FFAI/Models/Text/JambaText.swift b/Sources/FFAI/Models/Text/JambaText.swift
index fd609d97..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.
diff --git a/Sources/FFAI/Models/Text/LFM2Text.swift b/Sources/FFAI/Models/Text/LFM2Text.swift
index 66517bbc..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.
diff --git a/Sources/FFAI/Models/Text/Qwen3xText.swift b/Sources/FFAI/Models/Text/Qwen3xText.swift
index d32ba889..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.
diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift
index 02e887c5..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.
diff --git a/Sources/FFAI/Ops/OpsFused.swift b/Sources/FFAI/Ops/OpsFused.swift
index 808bc5cf..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.
diff --git a/Sources/FFAI/Ops/OpsValidation.swift b/Sources/FFAI/Ops/OpsValidation.swift
index ca31d1de..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.
diff --git a/Sources/MetalTileSwift/MetalTileLibrary.swift b/Sources/MetalTileSwift/MetalTileLibrary.swift
index a8e1bb89..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.
diff --git a/Tests/FFAITests/Benchmark/MoEBgemmBm64MppTests.swift b/Tests/FFAITests/Benchmark/MoEBgemmBm64MppTests.swift
index 8e254826..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.
diff --git a/Tests/FFAITests/Generation/DrafterTests.swift b/Tests/FFAITests/Generation/DrafterTests.swift
index 5aa380c0..75e6dac2 100644
--- a/Tests/FFAITests/Generation/DrafterTests.swift
+++ b/Tests/FFAITests/Generation/DrafterTests.swift
@@ -1,3 +1,17 @@
+// 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
diff --git a/Tests/FFAITests/Ops/OpsSpecialPathTests.swift b/Tests/FFAITests/Ops/OpsSpecialPathTests.swift
index 1f231248..3cecdc27 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.
diff --git a/Tests/FFAITests/Ops/OpsTests.swift b/Tests/FFAITests/Ops/OpsTests.swift
index cc30ca46..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.
diff --git a/Tests/FFAITests/Ops/OpsValidationTests.swift b/Tests/FFAITests/Ops/OpsValidationTests.swift
index ddb135f7..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.
diff --git a/Tests/FFAITests/Ops/QuantizedOpsTests.swift b/Tests/FFAITests/Ops/QuantizedOpsTests.swift
index 87818f21..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.
diff --git a/Tests/Helpers/CommonTestHelpers.swift b/Tests/Helpers/CommonTestHelpers.swift
index fe252742..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.
diff --git a/Tests/MetalTileSwiftTests/PSOCacheTests.swift b/Tests/MetalTileSwiftTests/PSOCacheTests.swift
index 973fbe22..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.
diff --git a/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift b/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift
index c60bd0c1..5b08a069 100644
--- a/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.swift
+++ b/Tests/ModelIntegrationTests/Text/Qwen35MoEBenchIntegrationTests.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.
diff --git a/Tests/ModelIntegrationTests/Text/Qwen36TextIntegrationTests.swift b/Tests/ModelIntegrationTests/Text/Qwen36TextIntegrationTests.swift
index f077ffee..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.
diff --git a/Tests/ModelIntegrationTests/Text/SpecDecodeBenchTests.swift b/Tests/ModelIntegrationTests/Text/SpecDecodeBenchTests.swift
index a105b74e..e524215a 100644
--- a/Tests/ModelIntegrationTests/Text/SpecDecodeBenchTests.swift
+++ b/Tests/ModelIntegrationTests/Text/SpecDecodeBenchTests.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.
diff --git a/Tests/ModelIntegrationTests/Text/SpecDecodeVerifyCostBench.swift b/Tests/ModelIntegrationTests/Text/SpecDecodeVerifyCostBench.swift
index 08406ef6..3f06325d 100644
--- a/Tests/ModelIntegrationTests/Text/SpecDecodeVerifyCostBench.swift
+++ b/Tests/ModelIntegrationTests/Text/SpecDecodeVerifyCostBench.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.

From 8c9e86302876cf8d17b0b3dee9671642c8ddb3fc Mon Sep 17 00:00:00 2001
From: Eric Kryski <599019+ekryski@users.noreply.github.com>
Date: Tue, 26 May 2026 23:56:39 -0600
Subject: [PATCH 54/54] test(ops): pin sdpaDecode2Pass against single-pass
 sdpaDecode
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

PR #10's regression was the wrapper accepting blocks=2 (kernel
pass-2 hardcodes a 32-lane block-merge reduction, so anything below
32 silently produces zero output). PR #14 closed it with a
`blocks >= 32 && blocks % 32 == 0` precondition.

The other half of the contract — that at the minimum legal
blocks=32, two-pass output matches the canonical single-pass kernel
within tolerance — had no host-side test. Add one. metaltile's per-
kernel GPU correctness test would not have caught the dispatch-
routing bug; the wrapper is what needs covering.
---
 Tests/FFAITests/Ops/OpsSpecialPathTests.swift | 70 +++++++++++++++++++
 1 file changed, 70 insertions(+)

diff --git a/Tests/FFAITests/Ops/OpsSpecialPathTests.swift b/Tests/FFAITests/Ops/OpsSpecialPathTests.swift
index 3cecdc27..d4faff49 100644
--- a/Tests/FFAITests/Ops/OpsSpecialPathTests.swift
+++ b/Tests/FFAITests/Ops/OpsSpecialPathTests.swift
@@ -471,6 +471,76 @@ struct OpsSpecialPathTests {
         }
     }
 
+    // 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")