diff --git a/Sources/FFAI/Models/MoELayer.swift b/Sources/FFAI/Models/MoELayer.swift index 80003eff..e3214a61 100644 --- a/Sources/FFAI/Models/MoELayer.swift +++ b/Sources/FFAI/Models/MoELayer.swift @@ -491,6 +491,10 @@ public final class MoELayer: Module, DecoderLayer { /// gate). Raw env string cached at init; `sortPlanCountingEnabled` /// below is the single place that interprets it. let sortPlanCountingEnvRaw: String? + /// Parallelized tile-plan builder gate (F-85 small-M follow-up). Raw + /// env string cached at init; `parallelPlanEnabled` below is the + /// single place that interprets it. + let parallelPlanEnvRaw: String? /// Lazy dequant of the gate weight + cached output. Skipped when /// `gate.inner` is not a 8-bit QuantizedLinear or env @@ -676,6 +680,7 @@ public final class MoELayer: Module, DecoderLayer { self.pairedIndirectEnvRaw = env["FFAI_MOE_PAIRED_INDIRECT"] self.coopIndirectEnvRaw = env["FFAI_MOE_COOP_INDIRECT"] self.sortPlanCountingEnvRaw = env["FFAI_MOE_COUNTING_SORT"] + self.parallelPlanEnvRaw = env["FFAI_MOE_PARALLEL_PLAN"] } /// F-85 expert-aligned tile-plan gate. Default ON as of the wiring @@ -712,6 +717,29 @@ public final class MoELayer: Module, DecoderLayer { sortPlanCountingEnvRaw != "0" } + /// F-85 parallelized tile-plan builder gate (small-M follow-up to the + /// idle-tile-cost/counting-sort round). `ffai_moe_build_tile_plan` / + /// `_bm32` / `_bm32_own` each dispatch ONE threadgroup for the whole + /// plan build - every expert's binary search, the prefix sum, and + /// every tile emission all run on a single GPU core, a per-layer fixed + /// cost that does not shrink with `mTotal` and was found + /// disproportionately expensive at low prefill `T`. The parallel + /// replacements (`Ops.moeBuildTilePlanParallel`/`Bm32Parallel`/ + /// `Bm32OwnParallel`, see `moe_tile_plan_expert_counts.rs`'s module + /// doc) split the per-expert row-range search out to `nExperts` + /// threadgroups (one per expert, running concurrently across the + /// GPU's cores) followed by a small single-threadgroup prefix-sum + + /// emission pass - same output format/contract, byte-exact against the + /// original builders across the full fixture set (dual CPU-oracle + + /// live-original-kernel checks). Default ON. Escape hatch: + /// `FFAI_MOE_PARALLEL_PLAN=0` reverts to the single-threadgroup + /// builders. Applies wherever `useTilePlan` holds - the plain bm16 + /// path, the default `useCoopTilePlan` path, and the paired lever's + /// "own" dispatch (see `decodeMany`). + public var parallelPlanEnabled: Bool { + parallelPlanEnvRaw != "0" + } + /// F-85 coop-core MoE gather GEMM gate (BM=32 tile plan over /// `ffai_moe_gather_qmm_coop`, the gather counterpart of the dense /// `ffai_qmm_coop` core). Default ON as of the wiring pass's quality + @@ -1461,11 +1489,23 @@ public final class MoELayer: Module, DecoderLayer { buffer: ownIndirectGateUpBuf, offset: 4, shape: [1], dtype: .u32) let ownTileCountDown = Tensor( buffer: ownIndirectDownBuf, offset: 4, shape: [1], dtype: .u32) - Ops.moeBuildTilePlanBm32Own( - sortedExperts: indices, - tileExpert: tileExpert, tileRowStart: tileRowStart, tileRowCount: tileRowCount, - tileCountGateUp: ownTileCountGateUp, tileCountDown: ownTileCountDown, - mTotal: mTotal, nExperts: nExperts, on: work) + // F-85 parallelized builder (see `parallelPlanEnabled`) - + // default ON, escape hatch FFAI_MOE_PARALLEL_PLAN=0. + if self.parallelPlanEnabled { + Ops.moeBuildTilePlanBm32OwnParallel( + sortedExperts: indices, + tileExpert: tileExpert, tileRowStart: tileRowStart, + tileRowCount: tileRowCount, + tileCountGateUp: ownTileCountGateUp, tileCountDown: ownTileCountDown, + mTotal: mTotal, nExperts: nExperts, on: work) + } else { + Ops.moeBuildTilePlanBm32Own( + sortedExperts: indices, + tileExpert: tileExpert, tileRowStart: tileRowStart, + tileRowCount: tileRowCount, + tileCountGateUp: ownTileCountGateUp, tileCountDown: ownTileCountDown, + mTotal: mTotal, nExperts: nExperts, on: work) + } coopIndirectArgsGateUp = ownIndirectGateUpBuf coopIndirectArgsDown = ownIndirectDownBuf @@ -1527,18 +1567,43 @@ public final class MoELayer: Module, DecoderLayer { buffer: indirectGateUpBuf, offset: 4, shape: [1], dtype: .u32) let tileCountDown = Tensor( buffer: indirectDownBuf, offset: 4, shape: [1], dtype: .u32) - Ops.moeBuildTilePlanBm32( - sortedExperts: indices, - tileExpert: tileExpert, tileRowStart: tileRowStart, tileRowCount: tileRowCount, - tileCountGateUp: tileCountGateUp, tileCountDown: tileCountDown, - mTotal: mTotal, nExperts: nExperts, on: work) + // F-85 parallelized builder (see `parallelPlanEnabled`) - + // default ON, escape hatch FFAI_MOE_PARALLEL_PLAN=0. This is + // the default production coop path's builder call. + if self.parallelPlanEnabled { + Ops.moeBuildTilePlanBm32Parallel( + sortedExperts: indices, + tileExpert: tileExpert, tileRowStart: tileRowStart, + tileRowCount: tileRowCount, + tileCountGateUp: tileCountGateUp, tileCountDown: tileCountDown, + mTotal: mTotal, nExperts: nExperts, on: work) + } else { + Ops.moeBuildTilePlanBm32( + sortedExperts: indices, + tileExpert: tileExpert, tileRowStart: tileRowStart, + tileRowCount: tileRowCount, + tileCountGateUp: tileCountGateUp, tileCountDown: tileCountDown, + mTotal: mTotal, nExperts: nExperts, on: work) + } coopIndirectArgsGateUp = indirectGateUpBuf coopIndirectArgsDown = indirectDownBuf } else { - Ops.moeBuildTilePlan( - sortedExperts: indices, - tileExpert: tileExpert, tileRowStart: tileRowStart, tileRowCount: tileRowCount, - mTotal: mTotal, nExperts: nExperts, on: work) + // F-85 parallelized builder (see `parallelPlanEnabled`) - + // bm16 fallback path (coop core off or shape doesn't + // qualify - see `useCoopTilePlan`). + if self.parallelPlanEnabled { + Ops.moeBuildTilePlanParallel( + sortedExperts: indices, + tileExpert: tileExpert, tileRowStart: tileRowStart, + tileRowCount: tileRowCount, + mTotal: mTotal, nExperts: nExperts, on: work) + } else { + Ops.moeBuildTilePlan( + sortedExperts: indices, + tileExpert: tileExpert, tileRowStart: tileRowStart, + tileRowCount: tileRowCount, + mTotal: mTotal, nExperts: nExperts, on: work) + } } // xGathered is already pre-gathered (step 8 above) - identity // row map, same helper the bm16 MPP path uses. diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift index 076e0695..112c73ee 100644 --- a/Sources/FFAI/Ops/Ops.swift +++ b/Sources/FFAI/Ops/Ops.swift @@ -6457,6 +6457,184 @@ public enum Ops { gridSize: grid, threadgroupSize: tg, on: cmd) } + /// Phase 1 of the parallelized tile-plan builders (F-85 small-M + /// follow-up, `ffai_moe_tile_plan_expert_counts`). Dispatches ONE + /// THREADGROUP PER EXPERT (`nExperts` groups, one thread each) so every + /// expert's binary search over `sortedExperts` runs concurrently across + /// the GPU's cores instead of serially on the single core the original + /// builders' one-threadgroup dispatch confines all `nExperts` lanes to. + /// Writes `expertRowBase[e]`/`expertCount[e]`, the same row-range values + /// any of the three single-threadgroup builders would compute locally + /// for expert `e` - shared by all three tile-height variants below, + /// since the row range a sorted-expert array implies for an expert does + /// not depend on the tile height it will be sliced into. + /// + /// Callers (`moeBuildTilePlanParallel`/`Bm32Parallel`/`Bm32OwnParallel` + /// below) allocate `expertRowBase`/`expertCount` as fresh per-call + /// scratch, matching `moeSortPlanCounting`'s `blockCounts`/ + /// `blockOffsets` convention - not cached, since `nExperts` is fixed + /// per model but this keeps every caller's allocation lifetime obvious. + public static func moeTilePlanExpertCounts( + sortedExperts: Tensor, + expertRowBase: Tensor, expertCount: Tensor, + mTotal: Int, nExperts: Int, + on cmd: MTLCommandBuffer + ) { + precondition( + sortedExperts.dtype == .u32, "moeTilePlanExpertCounts: sortedExperts must be u32") + precondition( + expertRowBase.dtype == .u32 && expertCount.dtype == .u32, + "moeTilePlanExpertCounts: expertRowBase/expertCount must be u32") + precondition( + nExperts <= 256, + "moeTilePlanExpertCounts: nExperts (\(nExperts)) exceeds the kernel's 256-expert scope" + ) + let grid = MTLSize(width: nExperts, height: 1, depth: 1) + let tg = MTLSize(width: 1, height: 1, depth: 1) + MetalTileKernels.ffai_moe_tile_plan_expert_counts_f32_threadgroups( + sorted_experts: sortedExperts.buffer, sorted_expertsOffset: sortedExperts.offset, + expert_row_base: expertRowBase.buffer, expert_row_baseOffset: expertRowBase.offset, + expert_count: expertCount.buffer, expert_countOffset: expertCount.offset, + m_total: UInt32(mTotal), + gridSize: grid, threadgroupSize: tg, on: cmd) + } + + /// Parallelized sibling of `moeBuildTilePlan` (F-85 small-M follow-up): + /// dispatches `moeTilePlanExpertCounts` (phase 1, per-expert parallel) + /// then `ffai_moe_build_tile_plan_parallel` (phase 2 - prefix sum + + /// tile emission only, no per-lane binary search left). Identical + /// output format/contract to `moeBuildTilePlan` - same + /// `tileExpert`/`tileRowStart`/`tileRowCount`, same worst-case capacity + /// and zero-fill padding convention. Gated by `FFAI_MOE_PARALLEL_PLAN` + /// in `MoELayer.swift`; `moeBuildTilePlan` remains the escape hatch. + public static func moeBuildTilePlanParallel( + sortedExperts: Tensor, + tileExpert: Tensor, tileRowStart: Tensor, tileRowCount: Tensor, + mTotal: Int, nExperts: Int, + on cmd: MTLCommandBuffer + ) { + precondition( + sortedExperts.dtype == .u32, "moeBuildTilePlanParallel: sortedExperts must be u32") + precondition( + tileExpert.dtype == .u32 && tileRowStart.dtype == .u32 && tileRowCount.dtype == .u32, + "moeBuildTilePlanParallel: tile_expert/tile_row_start/tile_row_count must be u32") + precondition( + nExperts <= 256, + "moeBuildTilePlanParallel: nExperts (\(nExperts)) exceeds the kernel's 256-lane threadgroup scratch capacity" + ) + let rowBaseBuf = device.makeBuffer(length: nExperts * 4) + let countBuf = device.makeBuffer(length: nExperts * 4) + let expertRowBase = Tensor(buffer: rowBaseBuf, offset: 0, shape: [nExperts], dtype: .u32) + let expertCount = Tensor(buffer: countBuf, offset: 0, shape: [nExperts], dtype: .u32) + moeTilePlanExpertCounts( + sortedExperts: sortedExperts, expertRowBase: expertRowBase, expertCount: expertCount, + mTotal: mTotal, nExperts: nExperts, on: cmd) + + let grid = MTLSize(width: 1, height: 1, depth: 1) + let tg = MTLSize(width: nExperts, height: 1, depth: 1) + MetalTileKernels.ffai_moe_build_tile_plan_parallel_f32_threadgroups( + expert_row_base: expertRowBase.buffer, expert_row_baseOffset: expertRowBase.offset, + expert_count: expertCount.buffer, expert_countOffset: expertCount.offset, + tile_expert: tileExpert.buffer, tile_expertOffset: tileExpert.offset, + tile_row_start: tileRowStart.buffer, tile_row_startOffset: tileRowStart.offset, + tile_row_count: tileRowCount.buffer, tile_row_countOffset: tileRowCount.offset, + n_experts: UInt32(nExperts), + gridSize: grid, threadgroupSize: tg, on: cmd) + } + + /// Parallelized sibling of `moeBuildTilePlanBm32` (F-85 small-M + /// follow-up) - same two-phase split as `moeBuildTilePlanParallel`, + /// BM=32 tile height, including the `tileCountGateUp`/`tileCountDown` + /// indirect-dispatch outputs. Gated by `FFAI_MOE_PARALLEL_PLAN`; + /// `moeBuildTilePlanBm32` remains the escape hatch. + public static func moeBuildTilePlanBm32Parallel( + sortedExperts: Tensor, + tileExpert: Tensor, tileRowStart: Tensor, tileRowCount: Tensor, + tileCountGateUp: Tensor, tileCountDown: Tensor, + mTotal: Int, nExperts: Int, + on cmd: MTLCommandBuffer + ) { + precondition( + sortedExperts.dtype == .u32, "moeBuildTilePlanBm32Parallel: sortedExperts must be u32") + precondition( + tileExpert.dtype == .u32 && tileRowStart.dtype == .u32 && tileRowCount.dtype == .u32, + "moeBuildTilePlanBm32Parallel: tile_expert/tile_row_start/tile_row_count must be u32") + precondition( + tileCountGateUp.dtype == .u32 && tileCountDown.dtype == .u32, + "moeBuildTilePlanBm32Parallel: tileCountGateUp/tileCountDown must be u32") + precondition( + nExperts <= 256, + "moeBuildTilePlanBm32Parallel: nExperts (\(nExperts)) exceeds the kernel's 256-lane threadgroup scratch capacity" + ) + let rowBaseBuf = device.makeBuffer(length: nExperts * 4) + let countBuf = device.makeBuffer(length: nExperts * 4) + let expertRowBase = Tensor(buffer: rowBaseBuf, offset: 0, shape: [nExperts], dtype: .u32) + let expertCount = Tensor(buffer: countBuf, offset: 0, shape: [nExperts], dtype: .u32) + moeTilePlanExpertCounts( + sortedExperts: sortedExperts, expertRowBase: expertRowBase, expertCount: expertCount, + mTotal: mTotal, nExperts: nExperts, on: cmd) + + let grid = MTLSize(width: 1, height: 1, depth: 1) + let tg = MTLSize(width: nExperts, height: 1, depth: 1) + MetalTileKernels.ffai_moe_build_tile_plan_bm32_parallel_f32_threadgroups( + expert_row_base: expertRowBase.buffer, expert_row_baseOffset: expertRowBase.offset, + expert_count: expertCount.buffer, expert_countOffset: expertCount.offset, + tile_expert: tileExpert.buffer, tile_expertOffset: tileExpert.offset, + tile_row_start: tileRowStart.buffer, tile_row_startOffset: tileRowStart.offset, + tile_row_count: tileRowCount.buffer, tile_row_countOffset: tileRowCount.offset, + tile_count_gateup: tileCountGateUp.buffer, tile_count_gateupOffset: tileCountGateUp.offset, + tile_count_down: tileCountDown.buffer, tile_count_downOffset: tileCountDown.offset, + n_experts: UInt32(nExperts), + gridSize: grid, threadgroupSize: tg, on: cmd) + } + + /// Parallelized sibling of `moeBuildTilePlanBm32Own` (F-85 small-M + /// follow-up) - same two-phase split, excludes `1..=16`-row pairing + /// candidates like the original "own" builder. Gated by + /// `FFAI_MOE_PARALLEL_PLAN`; `moeBuildTilePlanBm32Own` remains the + /// escape hatch. + public static func moeBuildTilePlanBm32OwnParallel( + sortedExperts: Tensor, + tileExpert: Tensor, tileRowStart: Tensor, tileRowCount: Tensor, + tileCountGateUp: Tensor, tileCountDown: Tensor, + mTotal: Int, nExperts: Int, + on cmd: MTLCommandBuffer + ) { + precondition( + sortedExperts.dtype == .u32, "moeBuildTilePlanBm32OwnParallel: sortedExperts must be u32" + ) + precondition( + tileExpert.dtype == .u32 && tileRowStart.dtype == .u32 && tileRowCount.dtype == .u32, + "moeBuildTilePlanBm32OwnParallel: tile_expert/tile_row_start/tile_row_count must be u32") + precondition( + tileCountGateUp.dtype == .u32 && tileCountDown.dtype == .u32, + "moeBuildTilePlanBm32OwnParallel: tileCountGateUp/tileCountDown must be u32") + precondition( + nExperts <= 256, + "moeBuildTilePlanBm32OwnParallel: nExperts (\(nExperts)) exceeds the kernel's 256-lane threadgroup scratch capacity" + ) + let rowBaseBuf = device.makeBuffer(length: nExperts * 4) + let countBuf = device.makeBuffer(length: nExperts * 4) + let expertRowBase = Tensor(buffer: rowBaseBuf, offset: 0, shape: [nExperts], dtype: .u32) + let expertCount = Tensor(buffer: countBuf, offset: 0, shape: [nExperts], dtype: .u32) + moeTilePlanExpertCounts( + sortedExperts: sortedExperts, expertRowBase: expertRowBase, expertCount: expertCount, + mTotal: mTotal, nExperts: nExperts, on: cmd) + + let grid = MTLSize(width: 1, height: 1, depth: 1) + let tg = MTLSize(width: nExperts, height: 1, depth: 1) + MetalTileKernels.ffai_moe_build_tile_plan_bm32_own_parallel_f32_threadgroups( + expert_row_base: expertRowBase.buffer, expert_row_baseOffset: expertRowBase.offset, + expert_count: expertCount.buffer, expert_countOffset: expertCount.offset, + tile_expert: tileExpert.buffer, tile_expertOffset: tileExpert.offset, + tile_row_start: tileRowStart.buffer, tile_row_startOffset: tileRowStart.offset, + tile_row_count: tileRowCount.buffer, tile_row_countOffset: tileRowCount.offset, + tile_count_gateup: tileCountGateUp.buffer, tile_count_gateupOffset: tileCountGateUp.offset, + tile_count_down: tileCountDown.buffer, tile_count_downOffset: tileCountDown.offset, + n_experts: UInt32(nExperts), + gridSize: grid, threadgroupSize: tg, on: cmd) + } + /// Device-side PAIR-ONLY BM=32 tile-plan builder for the F-85 /// tile-pairing lever (`ffai_moe_build_tile_plan_bm32_paired`). Reads /// the SAME `sortedExperts` array as `moeBuildTilePlanBm32Own` and diff --git a/Tests/ModelIntegrationTests/Text/TilePlanParallelGateTests.swift b/Tests/ModelIntegrationTests/Text/TilePlanParallelGateTests.swift new file mode 100644 index 00000000..bb305c28 --- /dev/null +++ b/Tests/ModelIntegrationTests/Text/TilePlanParallelGateTests.swift @@ -0,0 +1,171 @@ +// 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. +// +// F-85 parallelized tile-plan builder e2e gate - sweeps the batched +// forwardMany prefill path across the T ladder (512/1024/2048/4096/ +// 16384/32768), reporting median tok/s per T. Reads +// FFAI_MOE_PARALLEL_PLAN from the process environment ONCE (same as +// production `MoELayer` init) and runs ONE config per invocation - run +// this test TWICE, once with the default environment (parallel ON) and +// once with FFAI_MOE_PARALLEL_PLAN=0 set BEFORE `swift test` starts, to +// compare. A single process loading the SAME config twice (to A/B +// within one run) was tried and rejected: the second `Model.load` call +// started before the first model's Metal-resident buffers had actually +// been released, briefly holding two ~35B-parameter models resident and +// pushing this shared M5 Max to under 2GB free physical memory. Two +// separate process launches (this file's actual shape) guarantees full +// OS-level reclamation between configs - matches the established +// per-config-process convention already used for other F-85 gates +// (e.g. `FFAI_MOE_COUNTING_SORT=0 swift test --filter ...`). +// +// GATE_LADDER_MAX_T caps the ladder to `targetT <= GATE_LADDER_MAX_T` +// for incremental validation before committing wall-clock/memory to +// the largest sentinels (16384/32768). + +import Foundation +import Metal +import TestHelpers +import Testing + +@testable import FFAI + +private let gateLocalPath = "/Users/tom/models/Qwen3.6-35B-A3B-4bit" +private let gateCheckpointAvailable = + FileManager.default.fileExists(atPath: gateLocalPath) + +@Suite( + "F-85 Parallel Tile-Plan Builder Gate", .serialized, + .enabled( + if: gateCheckpointAvailable, + "requires a local checkpoint at \(gateLocalPath)") +) +struct TilePlanParallelGateTests { + + /// One (T, trialCount) sweep point in the ladder. + private struct Point { + let targetT: Int + let trials: Int + } + + private static let fullLadder: [Point] = [ + Point(targetT: 512, trials: 7), + Point(targetT: 1024, trials: 7), + Point(targetT: 2048, trials: 3), + Point(targetT: 4096, trials: 3), + Point(targetT: 16384, trials: 3), + Point(targetT: 32768, trials: 3), + ] + + /// GATE_LADDER_MAX_T / GATE_LADDER_MIN_T bound the ladder to + /// `minT <= targetT <= maxT`. Running one (T, trials) point per + /// process (min == max) is the safe default after repeated ladders + /// within one process were found to accumulate memory across T + /// values (each T's KV cache / activation buffers were not fully + /// released before the next T's allocation, independent of the + /// per-iteration `autoreleasepool` already in the trial loop below) - + /// a single loaded model swept multiple T values close to OOMed this + /// shared M5 Max twice. One T per process guarantees a clean process + /// exit reclaims everything before the next point runs. + private static var ladder: [Point] { + let env = ProcessInfo.processInfo.environment + let maxT = env["GATE_LADDER_MAX_T"].flatMap { Int($0) } ?? Int.max + let minT = env["GATE_LADDER_MIN_T"].flatMap { Int($0) } ?? 0 + // GATE_TRIALS_OVERRIDE caps trial count at the largest T values - + // fewer large-buffer allocate/free cycles lowers peak memory risk + // on this shared machine at the cost of a noisier median. + let trialsOverride = env["GATE_TRIALS_OVERRIDE"].flatMap { Int($0) } + return fullLadder.filter { $0.targetT >= minT && $0.targetT <= maxT } + .map { point in + guard let cap = trialsOverride, cap < point.trials else { return point } + return Point(targetT: point.targetT, trials: cap) + } + } + + @Test("F-85 parallel tile-plan builder - e2e prefill ladder (one config per process)") + func tilePlanParallelGateLadder() async throws { + let parallelPlanEnvRaw = ProcessInfo.processInfo.environment["FFAI_MOE_PARALLEL_PLAN"] + let label = + parallelPlanEnvRaw == "0" + ? "single-threadgroup(FFAI_MOE_PARALLEL_PLAN=0)" : "parallel(default)" + print("=== \(label) ===") + + var optsBuilder = LoadOptions() + optsBuilder.prewarm = false + let opts = optsBuilder + let m: Model = try await ModelLoadLock.shared.loadSerially { + try await Model.load(gateLocalPath, options: opts) + } + 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) + + for point in Self.ladder { + var encoded = seedEncoded + while encoded.count < point.targetT { + encoded.append(contentsOf: seedEncoded) + } + encoded = Array(encoded.prefix(point.targetT)) + let T = encoded.count + + // Warm up Metal PSO + first-token JIT before timing, same + // convention as `runForwardManyBench`. Each iteration's + // cache/command-buffer allocation and encoding is wrapped in + // `autoreleasepool` - a tight Swift-concurrency loop never + // gets the per-runloop-tick autorelease drain a normal event + // loop would, so the MTLBuffers `forwardMany` allocates while + // encoding (tile-plan scratch, gathered activations, etc.) + // accumulate across iterations instead of freeing each one - + // the documented MTLBuffer/autoreleasepool leak class, worse + // here than usual because this loop repeats at every T up to + // 32768. + for _ in 0 ..< 2 { + let warmCmd: MTLCommandBuffer = autoreleasepool { + let warmCaches = qwen.makeLayerCaches() + let warmCmd = Device.shared.makeCommandBuffer() + _ = qwen.forwardMany( + tokenIds: encoded, startPosition: 0, + caches: warmCaches, on: warmCmd, device: Device.shared) + warmCmd.commit() + return warmCmd + } + await warmCmd.awaitCompletion() + } + + var secs: [Double] = [] + for _ in 0 ..< point.trials { + var t0 = Date() + let cmd: MTLCommandBuffer = autoreleasepool { + let caches = qwen.makeLayerCaches() + let cmd = Device.shared.makeCommandBuffer() + t0 = Date() + _ = qwen.forwardMany( + tokenIds: encoded, startPosition: 0, + caches: caches, on: cmd, device: Device.shared) + cmd.commit() + return cmd + } + await cmd.awaitCompletion() + secs.append(Date().timeIntervalSince(t0)) + } + secs.sort() + let median = secs[secs.count / 2] + let tps = Double(T) / median + print( + "GATE \(label) T=\(T) trials=\(point.trials): runs=\(secs.map { String(format: "%.3f", $0) }) median=\(String(format: "%.3f", median))s tps=\(String(format: "%.2f", tps))" + ) + } + } +}