From ca51ade75f3b94d2f6eb49b0798847aee57d09e0 Mon Sep 17 00:00:00 2001 From: TheTom Date: Thu, 21 May 2026 11:07:07 -0500 Subject: [PATCH 001/120] feat(stubs+ops): rebuild against metaltile-emit naming overhaul Builds against metaltile-ffai feat/ffai-kernel-pack @ 0e19a8e where metaltile-emit standardized kernel names (mt_*/ffai_*/vector_*) and gained 16 new families. Ops.swift adds: - moeGatherDequantGemmInt4 wrapper (bm16 BGEMM, env-gated mpp/nax variant) - moeGatherDequantGemmInt4Bm8 wrapper (topK=8 decode path) - gatedDeltaChunk + sdpaPrefillMma wrappers (batched-T prefill ops) - gatedDeltaPrepStep wrapper (fused GDN prep + recurrence) - sink_end/window_start passed into 9 ffai_sdpa_decode dispatch sites (d64 / d256 / d512 across f32 / f16 / bf16); the d128 sites already passed them. Sliding-window-with-sink fast path is now uniformly wired across head-dim variants. MetalTileSwift/Stubs/FFAIStubs.swift (new file, 36 fatalError shims): - aura_encode_int{2,3,4,8}_* (12) - aura_dequant_rotated_int{2,3,4,8}_* (12) - ffai_gemm_* / ffai_rope_yarn_* / ffai_sdpa_multi_* / mt_rms_norm_wide_* (12) The corresponding Rust kernel sources don't exist in metaltile-ffai; each stub traps with a pointer back to metaltile task #93. Compile succeeds against the regenerated MetalTileKernels.swift bindings. --- Sources/FFAI/Ops.swift | 514 +++++++++++++++++++ Sources/MetalTileSwift/Stubs/FFAIStubs.swift | 414 +++++++++++++++ 2 files changed, 928 insertions(+) create mode 100644 Sources/MetalTileSwift/Stubs/FFAIStubs.swift diff --git a/Sources/FFAI/Ops.swift b/Sources/FFAI/Ops.swift index 721a63c6..d4f2fbff 100644 --- a/Sources/FFAI/Ops.swift +++ b/Sources/FFAI/Ops.swift @@ -1841,6 +1841,7 @@ public enum Ops { head_dim: UInt32(headDim), n_kv: UInt32(nKV), kv_stride: UInt32(kvStride), heads_per_group: UInt32(headsPerGroup), + sink_end: UInt32(sinkEnd), window_start: UInt32(windowStart), scale: scale, gridSize: grid, threadgroupSize: tg, on: cmd) case (64, .f16): @@ -1852,6 +1853,7 @@ public enum Ops { head_dim: UInt32(headDim), n_kv: UInt32(nKV), kv_stride: UInt32(kvStride), heads_per_group: UInt32(headsPerGroup), + sink_end: UInt32(sinkEnd), window_start: UInt32(windowStart), scale: scale, gridSize: grid, threadgroupSize: tg, on: cmd) case (64, .bf16): @@ -1863,6 +1865,7 @@ public enum Ops { head_dim: UInt32(headDim), n_kv: UInt32(nKV), kv_stride: UInt32(kvStride), heads_per_group: UInt32(headsPerGroup), + sink_end: UInt32(sinkEnd), window_start: UInt32(windowStart), scale: scale, gridSize: grid, threadgroupSize: tg, on: cmd) case (256, .f32): @@ -1874,6 +1877,7 @@ public enum Ops { head_dim: UInt32(headDim), n_kv: UInt32(nKV), kv_stride: UInt32(kvStride), heads_per_group: UInt32(headsPerGroup), + sink_end: UInt32(sinkEnd), window_start: UInt32(windowStart), scale: scale, gridSize: grid, threadgroupSize: tg, on: cmd) case (256, .f16): @@ -1885,6 +1889,7 @@ public enum Ops { head_dim: UInt32(headDim), n_kv: UInt32(nKV), kv_stride: UInt32(kvStride), heads_per_group: UInt32(headsPerGroup), + sink_end: UInt32(sinkEnd), window_start: UInt32(windowStart), scale: scale, gridSize: grid, threadgroupSize: tg, on: cmd) case (256, .bf16): @@ -1896,6 +1901,7 @@ public enum Ops { head_dim: UInt32(headDim), n_kv: UInt32(nKV), kv_stride: UInt32(kvStride), heads_per_group: UInt32(headsPerGroup), + sink_end: UInt32(sinkEnd), window_start: UInt32(windowStart), scale: scale, gridSize: grid, threadgroupSize: tg, on: cmd) case (512, .f32): @@ -1907,6 +1913,7 @@ public enum Ops { head_dim: UInt32(headDim), n_kv: UInt32(nKV), kv_stride: UInt32(kvStride), heads_per_group: UInt32(headsPerGroup), + sink_end: UInt32(sinkEnd), window_start: UInt32(windowStart), scale: scale, gridSize: grid, threadgroupSize: tg, on: cmd) case (512, .f16): @@ -1918,6 +1925,7 @@ public enum Ops { head_dim: UInt32(headDim), n_kv: UInt32(nKV), kv_stride: UInt32(kvStride), heads_per_group: UInt32(headsPerGroup), + sink_end: UInt32(sinkEnd), window_start: UInt32(windowStart), scale: scale, gridSize: grid, threadgroupSize: tg, on: cmd) case (512, .bf16): @@ -1929,6 +1937,7 @@ public enum Ops { head_dim: UInt32(headDim), n_kv: UInt32(nKV), kv_stride: UInt32(kvStride), heads_per_group: UInt32(headsPerGroup), + sink_end: UInt32(sinkEnd), window_start: UInt32(windowStart), scale: scale, gridSize: grid, threadgroupSize: tg, on: cmd) default: @@ -2305,4 +2314,509 @@ public enum Ops { } return result } + + // ─── MoE batched gather GEMM (int4) ────────────────────────────── + // + // Wraps `mt_moe_gather_qmm_mma_int4_bm16_*` — one kernel launch that + // processes `mTotal` rows of activations, each row tagged with the + // expert id it routes to via `indices[mTotal]`. Replaces a Python- + // style serial-expert loop (`topK` dispatches × 3 projections = 3K + // launches) with one batched dispatch per projection. + // + // Weight layout (mlx affine int4 stacked): + // `weight [E, N, K/8]` u32 packed, row-major within each expert. + // `scales`, `biases` `[E, N, K/groupSize]` in the activation dtype. + // Activation layout: + // `input [mTotal, K]` contiguous in `input.dtype`. + // `indices [mTotal]` u32 — expert id per row. ROWS MUST BE SORTED + // BY EXPERT (the kernel scans forward + // looking for an expert-id change to find + // the per-tile sub-range). + // Output: + // `out [mTotal, N]` in `input.dtype` (caller pre-allocates). + // + // Tile contract: N % 32 == 0, K % 32 == 0 (the kernel pads m_total + // internally — BM=16 with boundary masking). Grid = + // [N / 32, ceil(mTotal / 16), 1] + // Threadgroup = [64, 1, 1] (2 simdgroups). + public static func moeGatherDequantGemmInt4( + input: Tensor, + weight: Tensor, scales: Tensor, biases: Tensor, + indices: Tensor, + mTotal: Int, nOut: Int, kIn: Int, + groupSize: Int, + on cmd: MTLCommandBuffer, + into out: Tensor + ) { + // Shape / dtype invariants. The kernel itself does no validation + // so any mismatch silently produces garbage. + precondition(weight.dtype == .u32, "moeGatherDequantGemmInt4: weight must be u32 packed") + precondition(indices.dtype == .u32, + "moeGatherDequantGemmInt4: indices must be u32 (\(indices.dtype))") + precondition(input.dtype == scales.dtype && scales.dtype == biases.dtype, + "moeGatherDequantGemmInt4: input/scales/biases dtype must match") + precondition(out.dtype == input.dtype, + "moeGatherDequantGemmInt4: out dtype must match input") + precondition(input.elementCount == mTotal * kIn, + "moeGatherDequantGemmInt4: input has \(input.elementCount) elements, expected \(mTotal * kIn) (mTotal=\(mTotal) * kIn=\(kIn))") + precondition(out.elementCount == mTotal * nOut, + "moeGatherDequantGemmInt4: out has \(out.elementCount) elements, expected \(mTotal * nOut)") + precondition(indices.elementCount == mTotal, + "moeGatherDequantGemmInt4: indices has \(indices.elementCount) elements, expected mTotal=\(mTotal)") + precondition(nOut % 32 == 0, + "moeGatherDequantGemmInt4: nOut (\(nOut)) must be multiple of 32 for bm16 kernel") + precondition(kIn % 32 == 0, + "moeGatherDequantGemmInt4: kIn (\(kIn)) must be multiple of 32 for bm16 kernel") + + // BM=16, BN=32. The simdgroup-matrix variant uses 64 threads + // (2 SGs); the MPP / NAX cooperative-tensor variant uses 32 + // threads (1 SG). Canonical metaltile dispatch + // dispatch_with_grid(grid=[N/32, ceil(T/16), 1], tg=[TG_W, 1, 1]) + // calls `dispatchThreadgroups` (grid counted in TGs); the + // generated Swift wrapper calls `dispatchThreads` (grid counted + // in TOTAL threads) so total threads = grid × tg per axis. + // + // FFAI_MOE_BGEMM_MPP=1 opts into the NAX cooperative-tensor + // path — macOS 26.2+ on gen-≥17 GPU silicon only (M5 Max yes). + let useMpp = ProcessInfo.processInfo.environment["FFAI_MOE_BGEMM_MPP"] != nil + let tgWidth = useMpp ? 32 : 64 + let grid = MTLSize(width: (nOut / 32) * tgWidth, + height: (mTotal + 15) / 16, + depth: 1) + let tg = MTLSize(width: tgWidth, height: 1, depth: 1) + let mTotalU = UInt32(mTotal) + let nOutU = UInt32(nOut) + let kInU = UInt32(kIn) + let groupSizeU = UInt32(groupSize) + + if useMpp { + switch input.dtype { + case .f32: + MetalTileKernels.mt_moe_gather_qmm_mma_int4_bm16_mpp_f32( + x: input.buffer, xOffset: input.offset, + w: weight.buffer, wOffset: weight.offset, + scales: scales.buffer, scalesOffset: scales.offset, + biases: biases.buffer, biasesOffset: biases.offset, + indices: indices.buffer, indicesOffset: indices.offset, + out: out.buffer, outOffset: out.offset, + m_total: mTotalU, n_out: nOutU, k_in: kInU, + group_size: groupSizeU, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.mt_moe_gather_qmm_mma_int4_bm16_mpp_f16( + x: input.buffer, xOffset: input.offset, + w: weight.buffer, wOffset: weight.offset, + scales: scales.buffer, scalesOffset: scales.offset, + biases: biases.buffer, biasesOffset: biases.offset, + indices: indices.buffer, indicesOffset: indices.offset, + out: out.buffer, outOffset: out.offset, + m_total: mTotalU, n_out: nOutU, k_in: kInU, + group_size: groupSizeU, + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError("Ops.moeGatherDequantGemmInt4: MPP variant only emits f32/f16 (got \(input.dtype))") + } + return + } + + switch input.dtype { + case .f32: + MetalTileKernels.mt_moe_gather_qmm_mma_int4_bm16_f32( + x: input.buffer, xOffset: input.offset, + w: weight.buffer, wOffset: weight.offset, + scales: scales.buffer, scalesOffset: scales.offset, + biases: biases.buffer, biasesOffset: biases.offset, + indices: indices.buffer, indicesOffset: indices.offset, + out: out.buffer, outOffset: out.offset, + m_total: mTotalU, n_out: nOutU, k_in: kInU, + group_size: groupSizeU, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.mt_moe_gather_qmm_mma_int4_bm16_f16( + x: input.buffer, xOffset: input.offset, + w: weight.buffer, wOffset: weight.offset, + scales: scales.buffer, scalesOffset: scales.offset, + biases: biases.buffer, biasesOffset: biases.offset, + indices: indices.buffer, indicesOffset: indices.offset, + out: out.buffer, outOffset: out.offset, + m_total: mTotalU, n_out: nOutU, k_in: kInU, + group_size: groupSizeU, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.mt_moe_gather_qmm_mma_int4_bm16_bf16( + x: input.buffer, xOffset: input.offset, + w: weight.buffer, wOffset: weight.offset, + scales: scales.buffer, scalesOffset: scales.offset, + biases: biases.buffer, biasesOffset: biases.offset, + indices: indices.buffer, indicesOffset: indices.offset, + out: out.buffer, outOffset: out.offset, + m_total: mTotalU, n_out: nOutU, k_in: kInU, + group_size: groupSizeU, + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError("Ops.moeGatherDequantGemmInt4: unsupported dtype \(input.dtype)") + } + } + + // MARK: - Batched-T prefill kernels (Qwen3.5 / Qwen3.6 hybrid prefill) + + /// Multi-token Gated Delta Net recurrence over a chunk of `T` tokens. + /// + /// Wraps `mt_gated_delta_chunk` — same recurrence math as + /// `gatedDeltaStep` but runs the per-token loop *inside* the kernel + /// with the recurrent state kept in per-lane registers across the + /// entire `T` sweep. A single dispatch replaces `T` independent + /// `gatedDeltaStep` calls; the state buffer is read once at entry + /// and written once at exit. + /// + /// Layout (matches the MLX-LM `_make_gated_delta_kernel` convention + /// the metaltile-ffai kernel implements): + /// * `q, k` : `[T, Hk, Dk]` row-major (B=1) + /// * `v, y` : `[T, Hv, Dv]` row-major + /// * `g, beta` : `[T, Hv]` row-major + /// * `stateIn / stateOut` : `[Hv, Dv, Dk]` (one state per `hv`) + /// * `tLen` : `[1]` u32 — number of tokens in this chunk (runtime + /// scalar, NOT a constexpr; same PSO works for every + /// chunk length). + /// + /// All input tensors must share the activation dtype `T` — the + /// kernel is emitted in f32 / f16 / bf16 variants. For Qwen3.5 the + /// state is f32 (see `GDNStateCache.dtype`); pass q/k/v/g/beta as f32 + /// tensors as well. + /// + /// Dispatch: grid `(Dv, Hv)` threadgroups, 32 threads (one simdgroup) + /// per group — identical to `mt_gated_delta_step` apart from the + /// runtime `tLen`. `Dk % 32 == 0` invariant applies (max Dk = 256, so + /// `n_per_t = Dk/32 ≤ 8` register entries per lane). + public static func gatedDeltaChunk( + q: Tensor, k: Tensor, v: Tensor, g: Tensor, beta: Tensor, + stateIn: Tensor, into y: Tensor, stateOut: Tensor, + tLen: Tensor, + numKeyHeads: Int, numValueHeads: Int, + keyHeadDim: Int, valueHeadDim: Int, + on cmd: MTLCommandBuffer + ) { + if let reason = OpsValidation.validateGatedDeltaStep( + keyHeadDim: keyHeadDim, valueHeadDim: valueHeadDim, + numKeyHeads: numKeyHeads, numValueHeads: numValueHeads + ) { + preconditionFailure("Ops.gatedDeltaChunk: \(reason)") + } + precondition(q.dtype == k.dtype && k.dtype == v.dtype + && v.dtype == g.dtype && g.dtype == beta.dtype + && beta.dtype == stateIn.dtype && stateIn.dtype == stateOut.dtype + && stateOut.dtype == y.dtype, + "Ops.gatedDeltaChunk: every tensor must share dtype") + precondition(tLen.dtype == .u32 && tLen.elementCount == 1, + "Ops.gatedDeltaChunk: tLen must be a [1] u32 scalar buffer") + + let lanesPerGroup = 32 + let grid = MTLSize(width: valueHeadDim * lanesPerGroup, + height: numValueHeads, depth: 1) + let tg = MTLSize(width: lanesPerGroup, height: 1, depth: 1) + + switch q.dtype { + case .f32: + MetalTileKernels.mt_gated_delta_chunk_f32( + q: q.buffer, qOffset: q.offset, k: k.buffer, kOffset: k.offset, + v: v.buffer, vOffset: v.offset, g: g.buffer, gOffset: g.offset, + beta: beta.buffer, betaOffset: beta.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: UInt32(keyHeadDim), dv: UInt32(valueHeadDim), + hv: UInt32(numValueHeads), hk: UInt32(numKeyHeads), + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.mt_gated_delta_chunk_f16( + q: q.buffer, qOffset: q.offset, k: k.buffer, kOffset: k.offset, + v: v.buffer, vOffset: v.offset, g: g.buffer, gOffset: g.offset, + beta: beta.buffer, betaOffset: beta.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: UInt32(keyHeadDim), dv: UInt32(valueHeadDim), + hv: UInt32(numValueHeads), hk: UInt32(numKeyHeads), + gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.mt_gated_delta_chunk_bf16( + q: q.buffer, qOffset: q.offset, k: k.buffer, kOffset: k.offset, + v: v.buffer, vOffset: v.offset, g: g.buffer, gOffset: g.offset, + beta: beta.buffer, betaOffset: beta.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: UInt32(keyHeadDim), dv: UInt32(valueHeadDim), + hv: UInt32(numValueHeads), hk: UInt32(numKeyHeads), + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError("Ops.gatedDeltaChunk: unsupported dtype \(q.dtype)") + } + } + + /// Batched-Q causal SDPA prefill via the simdgroup-matrix MMA kernel. + /// Replaces `T` independent `sdpaDecode` calls (each O(K) per + /// q-head, T separate dispatches) with one dispatch that tiles the + /// `T × K` attention matrix 32×16 with simdgroup_matrix MMAs. + /// + /// Layout (per `mt_sdpa_prefill_mma` in metaltile-ffai): + /// * `q` : `[n_q_heads, q_len, head_dim]` row-major (B=1) + /// * `k,v` : `[n_kv_heads, k_len, head_dim]` row-major + /// * `out` : `[n_q_heads, q_len, head_dim]` + /// + /// `qLen` must be a multiple of 32 (BQ tile). Caller is responsible + /// for padding (or for splitting the prefill into a 32-aligned chunk + /// plus a per-token tail that falls back to `sdpaDecode`). + /// `kLen ≥ qLen` covers the causal mask; the kernel internally bounds + /// each row's K walk to its absolute position (assumes the last + /// `qLen` queries correspond to the last `qLen` K rows — i.e. the + /// queries are appended at the end of the K cache). + /// + /// Grid: `(qLen / 32, n_q_heads, 1)` threadgroups, 128 threads per + /// group (4 simdgroups). Reduction-mode kernel. + public static func sdpaPrefillMma( + q: Tensor, k: Tensor, v: Tensor, + nQHeads: Int, nKVHeads: Int, headDim: Int, + qLen: Int, kLen: Int, scale: Float, + on cmd: MTLCommandBuffer, + into out: Tensor? = nil + ) -> Tensor { + precondition(qLen > 0 && qLen % 32 == 0, + "Ops.sdpaPrefillMma: qLen \(qLen) must be a positive multiple of 32") + precondition(kLen >= qLen, + "Ops.sdpaPrefillMma: kLen \(kLen) must be >= qLen \(qLen)") + precondition(nQHeads % nKVHeads == 0, + "Ops.sdpaPrefillMma: nQHeads \(nQHeads) must be a multiple of nKVHeads \(nKVHeads)") + precondition(q.dtype == k.dtype && k.dtype == v.dtype, + "Ops.sdpaPrefillMma: q/k/v must share dtype") + let result = out ?? Tensor.empty(shape: [nQHeads, qLen, headDim], dtype: q.dtype) + // Grid: (qLen/BQ, n_q_heads, 1) threadgroups, 128 threads per + // group (4 simdgroups). dispatchThreads counts threads, so width + // = (qLen/32) * 128, height = n_q_heads. + let threadsPerGroup = 128 + let bq = 32 + let grid = MTLSize(width: (qLen / bq) * threadsPerGroup, + height: nQHeads, depth: 1) + let tg = MTLSize(width: threadsPerGroup, height: 1, depth: 1) + let gqaFactor = nQHeads / nKVHeads + + switch q.dtype { + case .f32: + MetalTileKernels.mt_sdpa_prefill_mma_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, + q_len: UInt32(qLen), k_len: UInt32(kLen), + gqa_factor: UInt32(gqaFactor), + n_q_heads: UInt32(nQHeads), n_kv_heads: UInt32(nKVHeads), + scale: scale, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.mt_sdpa_prefill_mma_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, + q_len: UInt32(qLen), k_len: UInt32(kLen), + gqa_factor: UInt32(gqaFactor), + n_q_heads: UInt32(nQHeads), n_kv_heads: UInt32(nKVHeads), + scale: scale, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.mt_sdpa_prefill_mma_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, + q_len: UInt32(qLen), k_len: UInt32(kLen), + gqa_factor: UInt32(gqaFactor), + n_q_heads: UInt32(nQHeads), n_kv_heads: UInt32(nKVHeads), + scale: scale, + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError("Ops.sdpaPrefillMma: unsupported dtype \(q.dtype)") + } + return result + } + + // ─── MoE gather BGEMM int4 — BM=8 MPP variant ──────────────────── + // + // Half-height counterpart of `moeGatherDequantGemmInt4` tuned for + // topK=8 decode where `mTotal = 8`. BM=16 wastes 50% of the tile + // rows on the trailing boundary; BM=8 fills the tile exactly. + // Uses the MPP destination-only-cooperative path, so the kernel + // descriptor (M=8, N=32, K=16) clears the simdgroup-scope + // constraint that the all-cooperative `matmul2d` would reject. + // Requires macOS 26.2+ / Apple10 GPU (M5 Max). + public static func moeGatherDequantGemmInt4Bm8( + input: Tensor, + weight: Tensor, scales: Tensor, biases: Tensor, + indices: Tensor, + mTotal: Int, nOut: Int, kIn: Int, + groupSize: Int, + on cmd: MTLCommandBuffer, + into out: Tensor + ) { + precondition(weight.dtype == .u32, + "moeGatherDequantGemmInt4Bm8: weight must be u32 packed") + precondition(indices.dtype == .u32, + "moeGatherDequantGemmInt4Bm8: indices must be u32") + precondition(input.dtype == scales.dtype && scales.dtype == biases.dtype, + "moeGatherDequantGemmInt4Bm8: input/scales/biases dtype must match") + precondition(out.dtype == input.dtype, + "moeGatherDequantGemmInt4Bm8: out dtype must match input") + precondition(input.elementCount == mTotal * kIn, + "moeGatherDequantGemmInt4Bm8: input elements \(input.elementCount) != mTotal*kIn \(mTotal*kIn)") + precondition(out.elementCount == mTotal * nOut, + "moeGatherDequantGemmInt4Bm8: out elements \(out.elementCount) != mTotal*nOut \(mTotal*nOut)") + precondition(indices.elementCount == mTotal, + "moeGatherDequantGemmInt4Bm8: indices elements \(indices.elementCount) != mTotal \(mTotal)") + precondition(nOut % 32 == 0, + "moeGatherDequantGemmInt4Bm8: nOut \(nOut) must be multiple of 32") + precondition(kIn % 32 == 0, + "moeGatherDequantGemmInt4Bm8: kIn \(kIn) must be multiple of 32") + + // 1 simdgroup per TG, BM=8 → grid Y = ceil(m/8). + let tgWidth = 32 + let grid = MTLSize(width: (nOut / 32) * tgWidth, + height: (mTotal + 7) / 8, + depth: 1) + let tg = MTLSize(width: tgWidth, height: 1, depth: 1) + let mTotalU = UInt32(mTotal) + let nOutU = UInt32(nOut) + let kInU = UInt32(kIn) + let groupSizeU = UInt32(groupSize) + + switch input.dtype { + case .f32: + MetalTileKernels.mt_moe_gather_qmm_mma_int4_bm8_mpp_f32( + x: input.buffer, xOffset: input.offset, + w: weight.buffer, wOffset: weight.offset, + scales: scales.buffer, scalesOffset: scales.offset, + biases: biases.buffer, biasesOffset: biases.offset, + indices: indices.buffer, indicesOffset: indices.offset, + out: out.buffer, outOffset: out.offset, + m_total: mTotalU, n_out: nOutU, k_in: kInU, + group_size: groupSizeU, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.mt_moe_gather_qmm_mma_int4_bm8_mpp_f16( + x: input.buffer, xOffset: input.offset, + w: weight.buffer, wOffset: weight.offset, + scales: scales.buffer, scalesOffset: scales.offset, + biases: biases.buffer, biasesOffset: biases.offset, + indices: indices.buffer, indicesOffset: indices.offset, + out: out.buffer, outOffset: out.offset, + m_total: mTotalU, n_out: nOutU, k_in: kInU, + group_size: groupSizeU, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.mt_moe_gather_qmm_mma_int4_bm8_mpp_bf16( + x: input.buffer, xOffset: input.offset, + w: weight.buffer, wOffset: weight.offset, + scales: scales.buffer, scalesOffset: scales.offset, + biases: biases.buffer, biasesOffset: biases.offset, + indices: indices.buffer, indicesOffset: indices.offset, + out: out.buffer, outOffset: out.offset, + m_total: mTotalU, n_out: nOutU, k_in: kInU, + group_size: groupSizeU, + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError("Ops.moeGatherDequantGemmInt4Bm8: unsupported dtype \(input.dtype)") + } + } + + // ─── Fused GDN prep + recurrence step ───────────────────────────── + // + // One dispatch absorbs the per-head q/k RMSNorm + g/beta math + + // the existing GDN recurrence step, collapsing the 3 host + // commit+wait pairs in Qwen35GDNMixer.forward down to 1. Same + // dispatch geometry as `mt_gated_delta_step`: + // grid = [dv, B·hv, 1] + // tg = [32, 1, 1] + // `convOut` layout: `[B, 2·Hk·Dk + Hv·Dv]` — q | k | v slabs. + // `qNormWeight` / `kNormWeight`: `[Hk·Dk]` (pass 1.0 × invKeyScale + // for the unweighted path). + public static func gatedDeltaPrepStep( + convOut: Tensor, + aLog: Tensor, dtBias: Tensor, + aRaw: Tensor, bRaw: Tensor, + qNormWeight: Tensor, kNormWeight: Tensor, + stateIn: Tensor, stateOut: Tensor, y: Tensor, + batchSize: Int, dk: Int, dv: Int, hv: Int, hk: Int, + on cmd: MTLCommandBuffer + ) { + precondition(convOut.dtype == aLog.dtype && aLog.dtype == dtBias.dtype + && dtBias.dtype == aRaw.dtype && aRaw.dtype == bRaw.dtype + && bRaw.dtype == qNormWeight.dtype && qNormWeight.dtype == kNormWeight.dtype + && kNormWeight.dtype == stateIn.dtype && stateIn.dtype == stateOut.dtype + && stateOut.dtype == y.dtype, + "Ops.gatedDeltaPrepStep: every tensor must share dtype") + precondition(dk % 32 == 0, + "Ops.gatedDeltaPrepStep: dk \(dk) must be multiple of 32") + precondition(dv % 32 == 0, + "Ops.gatedDeltaPrepStep: dv \(dv) must be multiple of 32") + precondition(hv % hk == 0, + "Ops.gatedDeltaPrepStep: hv \(hv) must be a multiple of hk \(hk) (GQA)") + + let grid = MTLSize(width: dv, + height: batchSize * hv, + depth: 1) + let tg = MTLSize(width: 32, height: 1, depth: 1) + let dkU = UInt32(dk) + let dvU = UInt32(dv) + let hvU = UInt32(hv) + let hkU = UInt32(hk) + + switch convOut.dtype { + case .f32: + MetalTileKernels.mt_gated_delta_prep_step_f32( + conv_out: convOut.buffer, conv_outOffset: convOut.offset, + a_log: aLog.buffer, a_logOffset: aLog.offset, + dt_bias: dtBias.buffer, dt_biasOffset: dtBias.offset, + a_raw: aRaw.buffer, a_rawOffset: aRaw.offset, + b_raw: bRaw.buffer, b_rawOffset: bRaw.offset, + q_norm_weight: qNormWeight.buffer, q_norm_weightOffset: qNormWeight.offset, + k_norm_weight: kNormWeight.buffer, k_norm_weightOffset: kNormWeight.offset, + state_in: stateIn.buffer, state_inOffset: stateIn.offset, + state_out: stateOut.buffer, state_outOffset: stateOut.offset, + y: y.buffer, yOffset: y.offset, + dk: dkU, dv: dvU, hv: hvU, hk: hkU, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .f16: + MetalTileKernels.mt_gated_delta_prep_step_f16( + conv_out: convOut.buffer, conv_outOffset: convOut.offset, + a_log: aLog.buffer, a_logOffset: aLog.offset, + dt_bias: dtBias.buffer, dt_biasOffset: dtBias.offset, + a_raw: aRaw.buffer, a_rawOffset: aRaw.offset, + b_raw: bRaw.buffer, b_rawOffset: bRaw.offset, + q_norm_weight: qNormWeight.buffer, q_norm_weightOffset: qNormWeight.offset, + k_norm_weight: kNormWeight.buffer, k_norm_weightOffset: kNormWeight.offset, + state_in: stateIn.buffer, state_inOffset: stateIn.offset, + state_out: stateOut.buffer, state_outOffset: stateOut.offset, + y: y.buffer, yOffset: y.offset, + dk: dkU, dv: dvU, hv: hvU, hk: hkU, + gridSize: grid, threadgroupSize: tg, on: cmd) + case .bf16: + MetalTileKernels.mt_gated_delta_prep_step_bf16( + conv_out: convOut.buffer, conv_outOffset: convOut.offset, + a_log: aLog.buffer, a_logOffset: aLog.offset, + dt_bias: dtBias.buffer, dt_biasOffset: dtBias.offset, + a_raw: aRaw.buffer, a_rawOffset: aRaw.offset, + b_raw: bRaw.buffer, b_rawOffset: bRaw.offset, + q_norm_weight: qNormWeight.buffer, q_norm_weightOffset: qNormWeight.offset, + k_norm_weight: kNormWeight.buffer, k_norm_weightOffset: kNormWeight.offset, + state_in: stateIn.buffer, state_inOffset: stateIn.offset, + state_out: stateOut.buffer, state_outOffset: stateOut.offset, + y: y.buffer, yOffset: y.offset, + dk: dkU, dv: dvU, hv: hvU, hk: hkU, + gridSize: grid, threadgroupSize: tg, on: cmd) + default: + fatalError("Ops.gatedDeltaPrepStep: unsupported dtype \(convOut.dtype)") + } + } } diff --git a/Sources/MetalTileSwift/Stubs/FFAIStubs.swift b/Sources/MetalTileSwift/Stubs/FFAIStubs.swift new file mode 100644 index 00000000..b173ce57 --- /dev/null +++ b/Sources/MetalTileSwift/Stubs/FFAIStubs.swift @@ -0,0 +1,414 @@ +import Metal + +// Stub dispatch functions for kernels that FFAI references but +// metaltile-emit does not currently emit (no Rust kernel source in +// metaltile-ffai). These let the FFAI package compile; calling any +// of them at runtime traps with a clear message pointing back to +// task #93 / the metaltile-ffai emitter. +// +// To remove a stub: (a) implement the kernel in +// `crates/metaltile-std/src/...`, (b) register it in +// `crates/metaltile-emit/src/main.rs::register_kernels`, (c) +// re-emit, then drop the matching stub func below. + +extension MetalTileKernels { + // MARK: - aura_dequant_rotated (codebook product-quantized dequant) + + @inline(never) private static func unimplemented(_ name: String) -> Never { + fatalError("\(name): kernel not currently emitted by metaltile-emit. See task #93 — needs Rust kernel source in crates/metaltile-std/ + registration in crates/metaltile-emit/src/main.rs.") + } + + public static func aura_dequant_rotated_int2_f32( + packed: MTLBuffer, packedOffset: Int = 0, + norms: MTLBuffer, normsOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, tokens: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_dequant_rotated_int2_f16( + packed: MTLBuffer, packedOffset: Int = 0, + norms: MTLBuffer, normsOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, tokens: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_dequant_rotated_int2_bf16( + packed: MTLBuffer, packedOffset: Int = 0, + norms: MTLBuffer, normsOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, tokens: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_dequant_rotated_int3_f32( + packed: MTLBuffer, packedOffset: Int = 0, + norms: MTLBuffer, normsOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, tokens: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_dequant_rotated_int3_f16( + packed: MTLBuffer, packedOffset: Int = 0, + norms: MTLBuffer, normsOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, tokens: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_dequant_rotated_int3_bf16( + packed: MTLBuffer, packedOffset: Int = 0, + norms: MTLBuffer, normsOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, tokens: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_dequant_rotated_int4_f32( + packed: MTLBuffer, packedOffset: Int = 0, + norms: MTLBuffer, normsOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, tokens: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_dequant_rotated_int4_f16( + packed: MTLBuffer, packedOffset: Int = 0, + norms: MTLBuffer, normsOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, tokens: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_dequant_rotated_int4_bf16( + packed: MTLBuffer, packedOffset: Int = 0, + norms: MTLBuffer, normsOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, tokens: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_dequant_rotated_int8_f32( + packed: MTLBuffer, packedOffset: Int = 0, + norms: MTLBuffer, normsOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, tokens: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_dequant_rotated_int8_f16( + packed: MTLBuffer, packedOffset: Int = 0, + norms: MTLBuffer, normsOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, tokens: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_dequant_rotated_int8_bf16( + packed: MTLBuffer, packedOffset: Int = 0, + norms: MTLBuffer, normsOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, tokens: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + // MARK: - aura_encode (codebook product-quantized encode) + + public static func aura_encode_int2_f32( + input: MTLBuffer, inputOffset: Int = 0, + rotation: MTLBuffer, rotationOffset: Int = 0, + boundaries: MTLBuffer, boundariesOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + packed_out: MTLBuffer, packed_outOffset: Int = 0, + norms_out: MTLBuffer, norms_outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_encode_int2_f16( + input: MTLBuffer, inputOffset: Int = 0, + rotation: MTLBuffer, rotationOffset: Int = 0, + boundaries: MTLBuffer, boundariesOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + packed_out: MTLBuffer, packed_outOffset: Int = 0, + norms_out: MTLBuffer, norms_outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_encode_int2_bf16( + input: MTLBuffer, inputOffset: Int = 0, + rotation: MTLBuffer, rotationOffset: Int = 0, + boundaries: MTLBuffer, boundariesOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + packed_out: MTLBuffer, packed_outOffset: Int = 0, + norms_out: MTLBuffer, norms_outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_encode_int3_f32( + input: MTLBuffer, inputOffset: Int = 0, + rotation: MTLBuffer, rotationOffset: Int = 0, + boundaries: MTLBuffer, boundariesOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + packed_out: MTLBuffer, packed_outOffset: Int = 0, + norms_out: MTLBuffer, norms_outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_encode_int3_f16( + input: MTLBuffer, inputOffset: Int = 0, + rotation: MTLBuffer, rotationOffset: Int = 0, + boundaries: MTLBuffer, boundariesOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + packed_out: MTLBuffer, packed_outOffset: Int = 0, + norms_out: MTLBuffer, norms_outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_encode_int3_bf16( + input: MTLBuffer, inputOffset: Int = 0, + rotation: MTLBuffer, rotationOffset: Int = 0, + boundaries: MTLBuffer, boundariesOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + packed_out: MTLBuffer, packed_outOffset: Int = 0, + norms_out: MTLBuffer, norms_outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_encode_int4_f32( + input: MTLBuffer, inputOffset: Int = 0, + rotation: MTLBuffer, rotationOffset: Int = 0, + boundaries: MTLBuffer, boundariesOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + packed_out: MTLBuffer, packed_outOffset: Int = 0, + norms_out: MTLBuffer, norms_outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_encode_int4_f16( + input: MTLBuffer, inputOffset: Int = 0, + rotation: MTLBuffer, rotationOffset: Int = 0, + boundaries: MTLBuffer, boundariesOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + packed_out: MTLBuffer, packed_outOffset: Int = 0, + norms_out: MTLBuffer, norms_outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_encode_int4_bf16( + input: MTLBuffer, inputOffset: Int = 0, + rotation: MTLBuffer, rotationOffset: Int = 0, + boundaries: MTLBuffer, boundariesOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + packed_out: MTLBuffer, packed_outOffset: Int = 0, + norms_out: MTLBuffer, norms_outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_encode_int8_f32( + input: MTLBuffer, inputOffset: Int = 0, + rotation: MTLBuffer, rotationOffset: Int = 0, + boundaries: MTLBuffer, boundariesOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + packed_out: MTLBuffer, packed_outOffset: Int = 0, + norms_out: MTLBuffer, norms_outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_encode_int8_f16( + input: MTLBuffer, inputOffset: Int = 0, + rotation: MTLBuffer, rotationOffset: Int = 0, + boundaries: MTLBuffer, boundariesOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + packed_out: MTLBuffer, packed_outOffset: Int = 0, + norms_out: MTLBuffer, norms_outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func aura_encode_int8_bf16( + input: MTLBuffer, inputOffset: Int = 0, + rotation: MTLBuffer, rotationOffset: Int = 0, + boundaries: MTLBuffer, boundariesOffset: Int = 0, + codebook: MTLBuffer, codebookOffset: Int = 0, + packed_out: MTLBuffer, packed_outOffset: Int = 0, + norms_out: MTLBuffer, norms_outOffset: Int = 0, + dim: UInt32, packed_width: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + // MARK: - ffai_gemm (naive matmul fallback) + + public static func ffai_gemm_f32( + weight: MTLBuffer, weightOffset: Int = 0, + input: MTLBuffer, inputOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + in_dim: UInt32, out_dim: UInt32, n_rows: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func ffai_gemm_f16( + weight: MTLBuffer, weightOffset: Int = 0, + input: MTLBuffer, inputOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + in_dim: UInt32, out_dim: UInt32, n_rows: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func ffai_gemm_bf16( + weight: MTLBuffer, weightOffset: Int = 0, + input: MTLBuffer, inputOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + in_dim: UInt32, out_dim: UInt32, n_rows: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + // MARK: - ffai_rope_yarn (YaRN-scaled RoPE) + + public static func ffai_rope_yarn_f32( + qk: MTLBuffer, qkOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + head_dim: UInt32, half_dim: UInt32, position: UInt32, + theta_base: Float, factor: Float, low: Float, high: Float, attn_factor: Float, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func ffai_rope_yarn_f16( + qk: MTLBuffer, qkOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + head_dim: UInt32, half_dim: UInt32, position: UInt32, + theta_base: Float, factor: Float, low: Float, high: Float, attn_factor: Float, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func ffai_rope_yarn_bf16( + qk: MTLBuffer, qkOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + head_dim: UInt32, half_dim: UInt32, position: UInt32, + theta_base: Float, factor: Float, low: Float, high: Float, attn_factor: Float, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + // MARK: - ffai_sdpa_multi (sparse-batched SDPA decode) + + public static func ffai_sdpa_multi_f32( + q: MTLBuffer, qOffset: Int = 0, + k: MTLBuffer, kOffset: Int = 0, + v: MTLBuffer, vOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + head_dim: UInt32, n_q_heads: UInt32, base_kv: UInt32, n_query: UInt32, + kv_stride: UInt32, heads_per_group: UInt32, + causal: UInt32, scale: Float, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func ffai_sdpa_multi_f16( + q: MTLBuffer, qOffset: Int = 0, + k: MTLBuffer, kOffset: Int = 0, + v: MTLBuffer, vOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + head_dim: UInt32, n_q_heads: UInt32, base_kv: UInt32, n_query: UInt32, + kv_stride: UInt32, heads_per_group: UInt32, + causal: UInt32, scale: Float, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func ffai_sdpa_multi_bf16( + q: MTLBuffer, qOffset: Int = 0, + k: MTLBuffer, kOffset: Int = 0, + v: MTLBuffer, vOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + head_dim: UInt32, n_q_heads: UInt32, base_kv: UInt32, n_query: UInt32, + kv_stride: UInt32, heads_per_group: UInt32, + causal: UInt32, scale: Float, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + // MARK: - mt_rms_norm_wide (wide-hidden RMSNorm) + + public static func mt_rms_norm_wide_f32( + x: MTLBuffer, xOffset: Int = 0, + w: MTLBuffer, wOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + eps_buf: MTLBuffer, eps_bufOffset: Int = 0, + n: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func mt_rms_norm_wide_f16( + x: MTLBuffer, xOffset: Int = 0, + w: MTLBuffer, wOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + eps_buf: MTLBuffer, eps_bufOffset: Int = 0, + n: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } + + public static func mt_rms_norm_wide_bf16( + x: MTLBuffer, xOffset: Int = 0, + w: MTLBuffer, wOffset: Int = 0, + out: MTLBuffer, outOffset: Int = 0, + eps_buf: MTLBuffer, eps_bufOffset: Int = 0, + n: UInt32, + gridSize: MTLSize, threadgroupSize: MTLSize, + on commandBuffer: MTLCommandBuffer + ) { unimplemented(#function) } +} From be2b85a0a4c88f7144909a6db69897f5a467885f Mon Sep 17 00:00:00 2001 From: TheTom Date: Thu, 21 May 2026 12:01:02 -0500 Subject: [PATCH 002/120] feat(ffai): wire bm8 BGEMM into MoELayer + add dequantGemmDynamicM driver MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two integration adds against metaltile-ffai feat/ffai-kernel-pack: * MoELayer.batchedSwiGLU: opt-in bm8 path via FFAI_MOE_BGEMM_BM8=1 when topK <= 8. Mirrors the existing FFAI_MOE_BGEMM (bm16) + FFAI_MOE_BGEMM_MPP (bm16_mpp) env gates. At T=1 Qwen3.6-A3B decode (mTotal=topK=8) the bm8 dispatch fills its BM=8 tile exactly — no trailing-row waste — but per benchShort smoke on M5 Max (16-step steady-state) the path is currently SLOWER than the sequential per-expert SwiGLU at decode (3.56 tps vs 4.80 baseline). The bottleneck at T=1 is gather-dispatch overhead, not tile-row utilization; bm8 is expected to win at prefill where mTotal scales with T. Left plumbed (off by default) so future bench cells can flip the gate. * Ops.dequantGemmDynamicM: host-side T-padding driver around mt_qmm_mma_{f32,f16,bf16}. Pads T -> ceil(T/32)*32 with zero rows, dispatches, slices the first T rows back. Fast-path is T already multiple of 32 (Qwen3.6 prefill chunk sizes 32/256/1024/4096/32768 are all naturally aligned). Bound to the new bf16 mt_qmm_mma binding emitted by metaltile-ffai eb091ef. Not wired into the forwardMany path yet — see /tmp/ffai-integration-roadmap.md for the decodeMany() refactor scope. Built clean: swift build -c release -p ffai. Qwen36SmokeTests/benchShort passes with FFAI_MOE_BGEMM=1 FFAI_MOE_BGEMM_BM8=1. --- Sources/FFAI/Models/MoELayer.swift | 223 +++++++++++++++++++++++++++-- Sources/FFAI/Ops.swift | 133 +++++++++++++++++ 2 files changed, 341 insertions(+), 15 deletions(-) diff --git a/Sources/FFAI/Models/MoELayer.swift b/Sources/FFAI/Models/MoELayer.swift index ebf4e343..eb183dc8 100644 --- a/Sources/FFAI/Models/MoELayer.swift +++ b/Sources/FFAI/Models/MoELayer.swift @@ -189,18 +189,76 @@ public final class MoELayer: Module, DecoderLayer { public let router: MoERouter public let hidden: Int + /// Stacked-expert weight handles for the batched gather BGEMM fast + /// path. When set, `decode` dispatches `mt_moe_gather_qmm_mma_int4_bm16` + /// instead of running `topK` sequential per-expert SwiGLU triplets — + /// 3 kernel launches per token instead of `3 * topK` (24 → 3 at + /// `topK=8`). Populated by the host model's MoE builder when (a) the + /// checkpoint is mlx int4-quantized, (b) the stacked weight tensors + /// are intact (not sliced away by the per-expert wrapper path), and + /// (c) the moeIntermediate / hidden shapes satisfy the bm16 kernel's + /// `N % 32 == 0` and `K % 32 == 0` tile contract. Left nil for + /// everyone else — the legacy serial-expert loop still runs. + public struct StackedInt4Experts: Sendable { + /// `[numExperts, moeIntermediate, hidden/8]` u32 packed. + public let gateWeight: Tensor + /// `[numExperts, moeIntermediate, hidden/groupSize]` in `dtype`. + public let gateScales: Tensor + public let gateBiases: Tensor + /// `[numExperts, moeIntermediate, hidden/8]` u32 packed. + public let upWeight: Tensor + public let upScales: Tensor + public let upBiases: Tensor + /// `[numExperts, hidden, moeIntermediate/8]` u32 packed. + public let downWeight: Tensor + public let downScales: Tensor + public let downBiases: Tensor + public let numExperts: Int + public let moeIntermediate: Int + public let hidden: Int + public let groupSize: Int + /// Activation dtype the scales / biases / activations all use. + public let dtype: DType + + public init(gateWeight: Tensor, gateScales: Tensor, gateBiases: Tensor, + upWeight: Tensor, upScales: Tensor, upBiases: Tensor, + downWeight: Tensor, downScales: Tensor, downBiases: Tensor, + numExperts: Int, moeIntermediate: Int, hidden: Int, + groupSize: Int, dtype: DType) { + self.gateWeight = gateWeight + self.gateScales = gateScales + self.gateBiases = gateBiases + self.upWeight = upWeight + self.upScales = upScales + self.upBiases = upBiases + self.downWeight = downWeight + self.downScales = downScales + self.downBiases = downBiases + self.numExperts = numExperts + self.moeIntermediate = moeIntermediate + self.hidden = hidden + self.groupSize = groupSize + self.dtype = dtype + } + } + public let stackedInt4Experts: StackedInt4Experts? + /// - gate: hidden → nExperts router projection. /// - gateProj/upProj/downProj: `nExperts`-long arrays of per-expert /// SwiGLU projections, index-aligned with the expert id. /// - sharedGate/Up/DownProj: optional shared-expert SwiGLU; pass all /// three or none. /// - router: the top-K + gating-math configuration. + /// - stackedInt4Experts: optional batched-BGEMM fast path. Per-expert + /// arrays above remain the source of truth for `parameters()` + /// (checkpoint binding) and the fallback decode path. public init(gate: AnyLinear, gateProj: [AnyLinear], upProj: [AnyLinear], downProj: [AnyLinear], sharedGateProj: AnyLinear? = nil, sharedUpProj: AnyLinear? = nil, sharedDownProj: AnyLinear? = nil, - router: MoERouter, hidden: Int) { + router: MoERouter, hidden: Int, + stackedInt4Experts: StackedInt4Experts? = nil) { precondition(gateProj.count == router.nExperts, "MoELayer: gateProj has \(gateProj.count) experts, router expects \(router.nExperts)") precondition(upProj.count == router.nExperts, @@ -220,6 +278,15 @@ public final class MoELayer: Module, DecoderLayer { self.sharedDownProj = sharedDownProj self.router = router self.hidden = hidden + if let s = stackedInt4Experts { + precondition(s.numExperts == router.nExperts, + "MoELayer: stackedInt4Experts.numExperts \(s.numExperts) ≠ router.nExperts \(router.nExperts)") + precondition(s.hidden == hidden, + "MoELayer: stackedInt4Experts.hidden \(s.hidden) ≠ MoELayer.hidden \(hidden)") + precondition(s.moeIntermediate % 32 == 0 && hidden % 32 == 0, + "MoELayer: stackedInt4Experts shape (moeIntermediate=\(s.moeIntermediate), hidden=\(hidden)) violates bm16 N%32 / K%32 tile contract") + } + self.stackedInt4Experts = stackedInt4Experts } public func parameters() -> [(String, Tensor)] { @@ -278,20 +345,42 @@ public final class MoELayer: Module, DecoderLayer { // contribute zero — skipping them is the same result, cheaper). let work = device.makeCommandBuffer() var accumulator: Tensor? - for (slot, expertId) in routing.indices.enumerated() { - // Broadcast the CPU combine weight into a [hidden] constant - // tensor so the element-wise `Ops.mul` can scale the expert - // output — avoids a dedicated scalar-multiply kernel. - let weightTensor = Tensor.filled(routing.weights[slot], - shape: [hidden], dtype: h.dtype, - device: device) - let expertOut = swiGLU(h, - gateProj: gateProj[expertId], - upProj: upProj[expertId], - downProj: downProj[expertId], - on: work) - let scaled = Ops.mul(expertOut, weightTensor, on: work) - accumulator = accumulator.map { Ops.add($0, scaled, on: work) } ?? scaled + // Batched gather-BGEMM fast path is opt-in via FFAI_MOE_BGEMM=1. + // At T=1 decode with topK=8 / m_total=8 it currently regresses + // vs the sequential per-expert matvec path on M5 Max — the + // bm16 tile pads to 16 rows but only 8 commit, so the kernel + // pays full weight-reload bandwidth for 50% useful work. Real + // win shape is prefill (m_total scales with T, more rows share + // each expert's weight tile) or NAX (FFAI_MOE_BGEMM_MPP=1). + // Path is wired + correctness-verified — flip the default once + // the m_total<16 regression is closed (likely via either a + // dedicated bm8 kernel emit or by promoting prefill to use + // this path first). + let enableBGEMM = ProcessInfo.processInfo.environment["FFAI_MOE_BGEMM"] != nil + if let stacked = stackedInt4Experts, stacked.dtype == h.dtype, enableBGEMM { + // Fast path: one batched gather BGEMM per projection. The + // kernel expects rows of activations sorted by expert id; we + // sort the topK indices ascending and replicate `h` into the + // gathered row order. The per-row expert assignment goes in + // an int32 indices buffer. + accumulator = batchedSwiGLU(h, stacked: stacked, routing: routing, + on: work, device: device) + } else { + for (slot, expertId) in routing.indices.enumerated() { + // Broadcast the CPU combine weight into a [hidden] constant + // tensor so the element-wise `Ops.mul` can scale the expert + // output — avoids a dedicated scalar-multiply kernel. + let weightTensor = Tensor.filled(routing.weights[slot], + shape: [hidden], dtype: h.dtype, + device: device) + let expertOut = swiGLU(h, + gateProj: gateProj[expertId], + upProj: upProj[expertId], + downProj: downProj[expertId], + on: work) + let scaled = Ops.mul(expertOut, weightTensor, on: work) + accumulator = accumulator.map { Ops.add($0, scaled, on: work) } ?? scaled + } } // ── 5. Optional always-on shared expert ────────────────────── @@ -307,6 +396,110 @@ public final class MoELayer: Module, DecoderLayer { return result } + /// Batched gather BGEMM fast path. T=1 decode: `mTotal = topK`. Each + /// row of the gathered batch holds the same `[hidden]` activation + /// vector but is paired with a different expert id via `indices`. + /// The kernel walks rows in tile order — boundary masking lets it + /// process the trailing partial tile when `mTotal < BM`. One + /// dispatch per projection replaces `topK` sequential per-expert + /// SwiGLU triplets: + /// 24 (topK=8) dispatches → 3 dispatches (gate / up / down) plus + /// element-wise SiLU + mul + scaled scatter-sum (~6 small ops). + /// Outputs are scaled by the router's combine weights on the CPU + /// (the same `Tensor.filled([hidden])` trick the serial path uses, + /// but only once per topK slot — still tiny vs the BGEMM cost). + /// + /// Tile selection: + /// - `FFAI_MOE_BGEMM_BM8=1` + `topK ≤ 8` → bm8 MPP kernel + /// (BM=8 fills the tile exactly at Qwen3.6-A3B topK=8 decode; + /// bm16 would waste 50 % of the trailing tile rows). + /// Requires Apple10+ GPU (M5 Max) + macOS 26.2+. + /// - otherwise → bm16 path (Ops.moeGatherDequantGemmInt4), + /// which itself respects `FFAI_MOE_BGEMM_MPP=1` for the + /// MPP/NAX variant at larger m_total. + private func batchedSwiGLU(_ h: Tensor, + stacked: StackedInt4Experts, + routing: MoERouter.Routing, + on cmd: MTLCommandBuffer, + device: Device) -> Tensor { + let topK = routing.indices.count + let moeIntermediate = stacked.moeIntermediate + let groupSize = stacked.groupSize + let dtype = stacked.dtype + let useBm8 = topK <= 8 + && ProcessInfo.processInfo.environment["FFAI_MOE_BGEMM_BM8"] != nil + + // ── Sort topK by expert id ascending (kernel contract) ── + // Track the original slot so we can apply the right combine + // weight after the BGEMM. The slot mapping never goes through + // the GPU, so sorting on the host is cheap. + let sorted = routing.indices.enumerated().sorted { $0.element < $1.element } + let sortedExperts = sorted.map { UInt32($0.element) } + let sortedSlots = sorted.map { $0.offset } + + // ── Materialise the [topK, hidden] activation gather + indices ── + let xGathered = Tensor.empty(shape: [topK, hidden], dtype: dtype, device: device) + let inputBytes = hidden * dtype.byteSize + // h is the same vector for every row at T=1 decode. memcpy from + // host backing — Tensors here are storage-shared. + let src = h.buffer.contents().advanced(by: h.offset) + let dst = xGathered.buffer.contents().advanced(by: xGathered.offset) + for r in 0..