diff --git a/Sources/FFAI/Models/MoELayer.swift b/Sources/FFAI/Models/MoELayer.swift index e3214a61..21313fed 100644 --- a/Sources/FFAI/Models/MoELayer.swift +++ b/Sources/FFAI/Models/MoELayer.swift @@ -605,6 +605,43 @@ public final class MoELayer: Module, DecoderLayer { private var coopIndirectArgsGateUpBuf: MTLBuffer? private var coopIndirectArgsDownBuf: MTLBuffer? + /// F-85 tile-plan buffers for the `useTilePlan` gather GEMM (plain + /// bm16 tileplan default AND the coop BM=32 default that supersedes + /// it, see `useTilePlan`/`useCoopTilePlan` in `decodeMany`) - allocated + /// lazily, per instance, then reused/grown across calls, unlike + /// `pairedTileExpertABuf` above this was still a fresh + /// `device.makeBuffer` triple every layer, every forward call on the + /// DEFAULT path (tilePlanEnabled/coopEnabled both default ON) until + /// this cache was added - the paired lever's own buffers got the + /// reuse treatment in the same round that shipped this gate, but the + /// plain/coop tile-plan buffers this gate always uses first were + /// missed. + /// + /// Unlike `pairedTileBufCapacity` (constant per instance - depends + /// only on `nExperts`), this capacity depends on `mTotal`, which + /// varies per prefill chunk (every chunk but the last is the same + /// configured step size, so in practice this grows once then holds + /// steady). Grow-only, matching `Ops.identityRowsBuffer`'s pattern: + /// a call needing less than the cached capacity just uses a + /// `[requestedCapacity]`-shaped view into the larger buffer (the + /// dispatch grid and the zero-fill below both use the PER-CALL + /// capacity, never the buffer's full backing length), so holding + /// onto the largest capacity seen so far is always safe and never + /// wastes a downstream dispatch on stale tail data. + /// + /// Same cross-call GPU-ordering safety argument as + /// `pairedTileExpertABuf` above (this file's whole per-layer + /// command-buffer chaining already depends on committed-order + /// execution with no inter-buffer fence), and the zero-fill every + /// call is still required for the same reason: a previous call's + /// tail past THIS call's real tile count would otherwise leak stale + /// `tile_row_count` values the GEMM kernel needs to read as zero to + /// treat that tile as inert padding. + private var tilePlanTileExpertBuf: MTLBuffer? + private var tilePlanTileRowStartBuf: MTLBuffer? + private var tilePlanTileRowCountBuf: MTLBuffer? + private var tilePlanBufCapacity = 0 + /// - gate: hidden → nExperts router projection. /// - gateProj/upProj/downProj: `nExperts`-long arrays of per-expert /// SwiGLU projections, index-aligned with the expert id. @@ -759,8 +796,24 @@ public final class MoELayer: Module, DecoderLayer { /// applies where `useTilePlan` already holds (DEVPLAN2 path) and the /// coop core's own shape gate passes (`nOut % 64 == 0` on both legs, /// `groupSize == 64`) - see `useCoopTilePlan` in `decodeMany`. + /// + /// Auto-default gates on `Ops.mppAutoCapable` (Apple GPU Family 9+ / + /// macOS 26+) - the same hardware requirement `dispatchQmmMma`'s coop + /// tier checks for the dense GEMM, since `ffai_moe_gather_qmm_coop` + /// wraps the identical `matmul2d` cooperative-tensor core. Unlike the + /// dense path (which falls back to a cheaper simdgroup tier when + /// incapable), this kernel has no such fallback once dispatched, so + /// the check has to live here rather than inside the Ops wrapper. + /// `FFAI_MOE_COOP=1` force-enables past this check (A/B only, same + /// "does not bypass the hardware requirement" caveat as + /// `FFAI_QMM_COOP=1` - forcing this on incapable hardware is + /// untested). public var coopEnabled: Bool { - coopEnvRaw != "0" + switch coopEnvRaw { + case "0": return false + case "1": return true + default: return Ops.mppAutoCapable + } } /// F-85 two-expert tile-pairing lever gate (`ffai_moe_gather_qmm_coop_paired`, @@ -912,6 +965,34 @@ public final class MoELayer: Module, DecoderLayer { return out } + /// Lazily allocates (grow-only, once-or-on-growth per instance) and + /// returns the three F-85 tile-plan buffers backing `useTilePlan` / + /// `useCoopTilePlan`, then zero-fills the requested `capacity` prefix + /// of each before handing them back. See the property doc above + /// `tilePlanTileExpertBuf` for the grow-only reuse design and why + /// re-zeroing the requested prefix every call is still required for + /// correctness (a previous, possibly larger, call's tail past THIS + /// call's real tile count would otherwise leak stale + /// expert/row_start/row_count into what the GEMM kernel expects to + /// be an inert dead tile). + private func tilePlanBuffers( + capacity: Int, device: Device + ) -> (tileExpert: MTLBuffer, tileRowStart: MTLBuffer, tileRowCount: MTLBuffer) { + if tilePlanTileExpertBuf == nil || tilePlanBufCapacity < capacity { + tilePlanTileExpertBuf = device.makeBuffer(length: capacity * 4) + tilePlanTileRowStartBuf = device.makeBuffer(length: capacity * 4) + tilePlanTileRowCountBuf = device.makeBuffer(length: capacity * 4) + tilePlanBufCapacity = capacity + } + let teBuf = tilePlanTileExpertBuf! + let trsBuf = tilePlanTileRowStartBuf! + let trcBuf = tilePlanTileRowCountBuf! + memset(teBuf.contents(), 0, capacity * 4) + memset(trsBuf.contents(), 0, capacity * 4) + memset(trcBuf.contents(), 0, capacity * 4) + return (teBuf, trsBuf, trcBuf) + } + /// Lazily allocates (once per instance) and returns the six F-85 /// paired-tile-plan buffers, then zero-fills all six before handing /// them back. See the property doc above `pairedTileExpertABuf` for @@ -1463,12 +1544,13 @@ public final class MoELayer: Module, DecoderLayer { useCoopTilePlan ? Ops.moeTilePlanCapacityBm32(mTotal: mTotal, nExperts: nExperts) : Ops.moeTilePlanCapacity(mTotal: mTotal, nExperts: nExperts) - let teBuf = device.makeBuffer(length: capacity * 4) - memset(teBuf.contents(), 0, capacity * 4) - let trsBuf = device.makeBuffer(length: capacity * 4) - memset(trsBuf.contents(), 0, capacity * 4) - let trcBuf = device.makeBuffer(length: capacity * 4) - memset(trcBuf.contents(), 0, capacity * 4) + // Preallocated, per-instance, grow-only reuse (see + // `tilePlanBuffers` / the property doc above + // `tilePlanTileExpertBuf`) - was three fresh `device.makeBuffer` + // calls every layer, every forward call on this gate's default + // (bm16-tileplan-or-coop) path, same allocation-churn class the + // paired lever's buffers were already fixed for. + let (teBuf, trsBuf, trcBuf) = tilePlanBuffers(capacity: capacity, device: device) let tileExpert = Tensor(buffer: teBuf, offset: 0, shape: [capacity], dtype: .u32) let tileRowStart = Tensor(buffer: trsBuf, offset: 0, shape: [capacity], dtype: .u32) let tileRowCount = Tensor(buffer: trcBuf, offset: 0, shape: [capacity], dtype: .u32) diff --git a/Sources/FFAI/Models/Text/Qwen3xText.swift b/Sources/FFAI/Models/Text/Qwen3xText.swift index 8407d562..5230d153 100644 --- a/Sources/FFAI/Models/Text/Qwen3xText.swift +++ b/Sources/FFAI/Models/Text/Qwen3xText.swift @@ -1265,6 +1265,11 @@ public final class Qwen35GDNLayerCache: LayerCacheProtocol, @unchecked Sendable // Force the coop kernel on even off the auto-capability check (A/B // only; the kernel needs the same MPP/NAX hardware support as // `ffai_qmm_mma_mpp`, this does not bypass that requirement). +// CAUTION: unlike the shape/group-size checks (still enforced), +// forcing this on hardware below the MPP/NAX capability floor is +// untested territory - the auto-capability gate exists specifically +// to keep this kernel off such hardware, and no test in this repo +// exercises the forced-on + incapable-hardware combination. // FFAI_QMM_MPP=0 | =1 // Force off / on the BN=32 MPP/NAX cooperative-tensor kernel // (`ffai_qmm_mma_mpp`), the tier below coop. Default: auto-detected @@ -1364,30 +1369,91 @@ public final class Qwen35GDNLayerCache: LayerCacheProtocol, @unchecked Sendable // at every T from 512 through 32768, no high-T regression. Tests: // kernels repo `moe_build_tile_plan_correctness.rs` + // `moe_gather_qmm_tileplan_correctness.rs`. -// FFAI_MOE_PAIRED=0 -// Opt out of the F-85 two-expert tile-pairing lever -// (`ffai_moe_gather_qmm_coop_paired`) at prefill, reverting to the -// coop core's single unpaired dispatch across every mTotal. Below +// FFAI_MOE_COOP=0 +// Opt out of the F-85 coop-core MoE gather GEMM (BM=32 tile plan +// over `ffai_moe_gather_qmm_coop`) at prefill, reverting to the +// BM=16 tileplan default. Only takes effect where `FFAI_MOE_TILEPLAN` +// is also active and the coop core's own shape gate passes +// (`nOut % 64 == 0` on both legs, `groupSize == 64`) - a layer that +// fails either check stays on the BM=16 tileplan path regardless of +// this flag. Default: auto-detected on `Ops.mppAutoCapable` hardware +// (Apple GPU Family 9+ / macOS 26+ - the same requirement +// `FFAI_QMM_COOP`'s auto-default checks, since this kernel wraps the +// same `matmul2d` cooperative-tensor core). `FFAI_MOE_COOP=1` force- +// enables past the hardware check (untested on incapable hardware, +// same caveat as `FFAI_QMM_COOP=1`). Isolated kernel bench shows +// +31-98% GB/s vs the BM=16 tileplan incumbent across mTotal +// 4096/16384/32768; `--dump-prefill-logits` bit-exact (cosine 1.0, +// max_abs_dlogit 0.0) at T=512/1024/4096/16384; e2e prefill win +// 14-23% at every T from 512 through 32768 (measured on +// `Ops.mppAutoCapable` hardware). +// FFAI_MOE_COOP_INDIRECT=0 +// Opt out of the idle-tile-cost indirect dispatch for the coop GEMM +// above (`dispatchThreadgroupsIndirect` fed the device-computed real +// tile count instead of the static worst-case-capacity bound), +// reverting to the old direct/capacity dispatch - kept for A/B +// measurement, not a production fallback (the old path pays 10-21% +// idle-tile cost at the F-85 low-context mTotal targets). Only takes +// effect where `FFAI_MOE_COOP` is also active (both the plain path +// and the paired lever's "own" dispatch below). Default: ON. +// FFAI_MOE_PAIRED=1 +// Opt IN to the F-85 two-expert tile-pairing lever +// (`ffai_moe_gather_qmm_coop_paired`) at prefill. Below // `FFAI_MOE_PAIRED_MAX_MTOTAL` (see below), the coop tile plan splits // into an "own" plan (full/17-31-row remainders, fed the UNCHANGED // `ffai_moe_gather_qmm_coop` dispatch) and a "paired" plan (1-16-row // remainders packed two experts to a tile, fed the new paired-GEMM // dispatch), eliminating the masked-out padding those short // remainders otherwise cost. Only takes effect where -// `FFAI_MOE_TILEPLAN` + `FFAI_MOE_COOP` are also active. Default: ON. -// Gated on `--dump-prefill-logits` bit-exact (cosine 1.0, -// max_abs_dlogit 0.0) + greedy-token match at T=512/1024/4096/16384; -// isolated kernel bench GO/NO-GO shows a 7-12% win at mTotal=4096, -// 4-9% at 8192, parity at 16384, and a 3-6% loss at 32768. Tests: -// kernels repo `moe_build_tile_plan_bm32_own_correctness.rs` + +// `FFAI_MOE_TILEPLAN` + `FFAI_MOE_COOP` are also active. +// DEFAULT OFF (opt-in, unlike every other MoE gate in this section - +// unset or `=0` stays off, `=1` enables). Isolated kernel bench +// GO/NO-GO shows a 7-12% win at mTotal=4096, 4-9% at 8192, parity at +// 16384, and a 3-6% loss at 32768, but the real wiring must always +// dispatch the paired GEMM at its static worst-case tile capacity +// (no GPU-to-CPU readback of the real paired-tile count before +// dispatch), which the indirect-dispatch fix below only partially +// recovers - see `pairedEnabled`'s doc in `MoELayer.swift` for the +// full measurement history of why this stays opt-in. Gated on +// `--dump-prefill-logits` bit-exact (cosine 1.0, max_abs_dlogit 0.0) +// + greedy-token match at T=512/1024/4096/16384. Tests: kernels repo +// `moe_build_tile_plan_bm32_own_correctness.rs` + // `moe_build_tile_plan_bm32_paired_correctness.rs` + // `moe_gather_qmm_coop_paired.rs`'s kernel_tests. +// FFAI_MOE_PAIRED_INDIRECT=0 +// Opt out of the paired-GEMM's own idle-tile-cost indirect dispatch +// (the root-cause fix for the worst-case-capacity regression +// `FFAI_MOE_PAIRED` above pays), reverting to the old direct/capacity +// dispatch - kept for A/B measurement against the indirect fix, not +// a production fallback. Only takes effect where `FFAI_MOE_PAIRED` +// is also active. Default: ON. // FFAI_MOE_PAIRED_MAX_MTOTAL= // Override the tile-pairing mTotal threshold (paired dispatch runs // when mTotal < this value). Default: 16384, the measured // parity boundary - chosen over 32768 (a measured 3-6% loss there) // so the default never dispatches the extra pass where it cannot // win. +// FFAI_MOE_COUNTING_SORT=0 +// Opt out of the F-85 bucketed counting-sort replacement for +// `ffai_moe_sort_plan` (histogram -> prefix-sum -> scatter, same +// stable expert-ascending order), reverting to the original kernel's +// O(mTotal^2) per-thread linear scan. Only takes effect on the +// DEVPLAN2 device-sorted-plan path (`useDevPlan2`) - the only path +// that calls the sort-plan kernel at all. Default: ON. +// `--dump-prefill-logits` byte-identical at every T tested; isolated +// kernel bench shows the win growing with mTotal: ~6-9x at +// 4096/8192, ~10-15x at 16384, ~38x at 32768. +// FFAI_MOE_PARALLEL_PLAN=0 +// Opt out of the parallelized MoE tile-plan builders +// (`ffai_moe_build_tile_plan_parallel` / `_bm32_parallel` / +// `_bm32_own_parallel` - one threadgroup per expert instead of a +// single threadgroup owning the whole plan build), reverting to the +// single-threadgroup builders. Applies wherever `FFAI_MOE_TILEPLAN` +// is active - the plain bm16 path, the `FFAI_MOE_COOP` path, and the +// `FFAI_MOE_PAIRED` lever's "own" dispatch. Default: ON. Byte-exact +// against the original single-threadgroup builders across the full +// kernels-repo fixture set (dual CPU-oracle + live-original-kernel +// checks). // DEVPLAN2=0 // Opt out of chaining the MoE prefill expert path onto the // caller's own command buffer. Note: no `FFAI_` prefix (older @@ -1462,6 +1528,19 @@ public final class Qwen35GDNMixer: Module { /// add + FFN chain onto the same command buffer. let fused: Bool + /// Chunked-WY prefill pipeline eligibility (shape half of `useWY` in + /// `forwardManyChunked`; the `t >= 64` half stays a per-call check + /// since `t` varies). Cached at init for the same reason as `fused` + /// above - `FFAI_GDN_WY` and the head-dim shape guards are all + /// instance-constant, so re-reading the env on every chunk was a + /// dictionary lookup this mixer's forward path never needed. + let wyShapeEligible: Bool + /// Shape-specialized fast prep_chunk kernel eligibility (shape + + /// env half of `useFast` in `forwardManyChunked`; the `!useWY` half + /// stays a per-call check since `useWY` depends on per-call `t`). + /// Cached at init, same reasoning as `wyShapeEligible` above. + let fastPrepChunkEligible: Bool + /// Pre-allocated per-call scratch tensors. The fused GDN path /// writes / reads these inside one command buffer per decode token; /// the engine's `workCmd.commit()` + caller wait between tokens @@ -1553,6 +1632,15 @@ public final class Qwen35GDNMixer: Module { // comparison. self.fused = ProcessInfo.processInfo.environment["FFAI_GDN_NO_FUSED_PREP"] == nil + // See `wyShapeEligible`/`fastPrepChunkEligible` doc above - cached + // once here instead of read from `forwardManyChunked` on every + // chunk. + self.wyShapeEligible = + ProcessInfo.processInfo.environment["FFAI_GDN_WY"] != nil + && valueHeadDim % 32 == 0 && keyHeadDim % 16 == 0 && valueHeadDim % 16 == 0 + self.fastPrepChunkEligible = + keyHeadDim == 128 && valueHeadDim == 128 && numValueHeads == 32 && numKeyHeads == 16 + && ProcessInfo.processInfo.environment["FFAI_GDN_NO_PREP_CHUNK_FAST"] == nil // Per-decode-token scratch — pre-allocated once at init so the // fused GDN path doesn't pay 6 × MTLBuffer allocations per call. @@ -2142,9 +2230,9 @@ public final class Qwen35GDNMixer: Module { // its doc comment), so the only remaining fallback to // `gatedDeltaPrepChunk` is `t < 64` (no full chunk to pad up to at // all). - let useWY = - ProcessInfo.processInfo.environment["FFAI_GDN_WY"] != nil && t >= 64 - && valueHeadDim % 32 == 0 && keyHeadDim % 16 == 0 && valueHeadDim % 16 == 0 + // Env + shape eligibility cached at init (`wyShapeEligible`) - only + // the `t >= 64` half varies per call. + let useWY = wyShapeEligible && t >= 64 // Shape-specialized fast path: dispatches // `ffai_gated_delta_prep_chunk_fast_d128_128_32_16` in place of the // generic `gatedDeltaPrepChunk`: same math, same grid/TG geometry, @@ -2155,12 +2243,10 @@ public final class Qwen35GDNMixer: Module { // doc for the root-cause writeup). Only valid at Qwen3.6-35B-A3B's // production GDN shape, so guard on the exact dims before dispatching, // and keep an escape hatch back to the generic kernel for any other - // GDN-bearing shape or for A/B verification. - let fastShapeMatches = - keyHeadDim == 128 && valueHeadDim == 128 && numValueHeads == 32 && numKeyHeads == 16 - let useFast = - !useWY && fastShapeMatches - && ProcessInfo.processInfo.environment["FFAI_GDN_NO_PREP_CHUNK_FAST"] == nil + // GDN-bearing shape or for A/B verification. Env + shape eligibility + // cached at init (`fastPrepChunkEligible`) - only the `!useWY` half + // varies per call (`useWY` itself depends on per-call `t`). + let useFast = !useWY && fastPrepChunkEligible if useWY { Ops.gatedDeltaWYPrefill( convOut: convOutAllF32, @@ -2291,6 +2377,18 @@ public final class Qwen35AttentionMixer: Module { return 1024 // default: only kick in at long KV }() + /// D=256 attention kernel tier, cached at init (same reasoning as + /// `sdpa2PassThreshold` above): `headDim` is instance-constant, so + /// re-deriving this from `ProcessInfo.environment` on every + /// `forwardMany` call (once per attention layer per prefill chunk, + /// and once per attention layer per decode token when the batched + /// decode path is active - see `FFAI_LEGACY_FORWARDMANY`) always + /// produced the same answer for the life of the mixer. Set together + /// in `init` below since `useLegacyD256Cached` only feeds + /// `useMmaD256Cached`'s derivation and has no other caller. + private let useMmaD256Cached: Bool + private let useQtiledD256Cached: Bool + init( qProj: AnyLinear, kProj: AnyLinear, vProj: AnyLinear, oProj: AnyLinear, qNorm: RMSNorm, kNorm: RMSNorm, @@ -2311,6 +2409,16 @@ public final class Qwen35AttentionMixer: Module { self.ropeTheta = ropeTheta self.attnOutputGate = attnOutputGate self.scale = 1.0 / Float(Double(headDim).squareRoot()) + // See `useMmaD256Cached`/`useQtiledD256Cached` doc above - cached + // once here instead of read from `forwardMany` on every call. + let qtiledD256 = + headDim == 256 + && ProcessInfo.processInfo.environment["FFAI_SDPA_QTILED_D256"] == "1" + let legacyD256 = + headDim == 256 && !qtiledD256 + && ProcessInfo.processInfo.environment["FFAI_SDPA_LEGACY_D256"] == "1" + self.useQtiledD256Cached = qtiledD256 + self.useMmaD256Cached = headDim == 256 && !legacyD256 && !qtiledD256 if attnOutputGate { // Pre-build the row-index tensors for the gate-split // gather. Even rows = queries, odd rows = gates. @@ -2732,14 +2840,15 @@ public final class Qwen35AttentionMixer: Module { // `ffai_sdpa_multi_d256` (one threadgroup per (query, q_head), // no K/V reuse across rows) that the F-85 campaign identified // as the source of the original append-regime bandwidth waste. - let useQtiledD256 = - headDim == 256 - && ProcessInfo.processInfo.environment["FFAI_SDPA_QTILED_D256"] == "1" - let useLegacyD256 = - headDim == 256 && !useQtiledD256 - && ProcessInfo.processInfo.environment["FFAI_SDPA_LEGACY_D256"] == "1" - let useMmaD256 = headDim == 256 && !useLegacyD256 && !useQtiledD256 - let useQtiled = useQtiledD256 + // Tier selection is instance-constant (depends only on `headDim` + // and the env, both fixed for this mixer's lifetime) - cached + // once in `init` as `useMmaD256Cached`/`useQtiledD256Cached` + // instead of re-read from `ProcessInfo.environment` on every + // `forwardMany` call (once per attention layer per prefill chunk, + // and once per attention layer per decode token on the batched + // decode path). + let useMmaD256 = useMmaD256Cached + let useQtiled = useQtiledD256Cached let (cacheK, cacheV) = kv.prepareForAttention(on: cmd) let attnAll: Tensor if useMmaD256 { diff --git a/Sources/FFAI/Ops/Ops.swift b/Sources/FFAI/Ops/Ops.swift index 112c73ee..4489553b 100644 --- a/Sources/FFAI/Ops/Ops.swift +++ b/Sources/FFAI/Ops/Ops.swift @@ -4574,7 +4574,15 @@ public enum Ops { /// load across twice as many query rows, at a threadgroup-memory /// cost that only fits f16/bf16 (33792 B at f32 would blow the 32 /// KiB ceiling, so f32 always stays on the BK=8 kernel below). - /// `FFAI_SDPA_BK16=0` reverts f16/bf16 to BK=8 for bisection. + /// `FFAI_SDPA_BK16=0` reverts f16/bf16 to BK=8 for bisection. Cached + /// once (`sdpaBK16Enabled` below) rather than read from + /// `ProcessInfo.environment` on every call - this is the default + /// d=256 prefill attention kernel, dispatched once per attention + /// layer per prefill chunk and once per attention layer per decode + /// token on the batched decode path. + private static let sdpaBK16Enabled: Bool = + ProcessInfo.processInfo.environment["FFAI_SDPA_BK16"] != "0" + public static func sdpaPrefillMmaD256( q: Tensor, k: Tensor, v: Tensor, nQHeads: Int, nKVHeads: Int, headDim: Int, @@ -4603,7 +4611,7 @@ public enum Ops { let tgGrid = MTLSize(width: tgX, height: nQHeads, depth: 1) let tg = MTLSize(width: 128, height: 1, depth: 1) let causalFlag = UInt32(causal ? 1 : 0) - let useBK16 = ProcessInfo.processInfo.environment["FFAI_SDPA_BK16"] != "0" + let useBK16 = Self.sdpaBK16Enabled switch q.dtype { case .f32: MetalTileKernels.ffai_sdpa_prefill_mma_d256_f32_threadgroups( @@ -5820,7 +5828,17 @@ public enum Ops { /// Whether the MPP auto-default may fire: Apple GPU Family 9+ (M3+) /// and macOS 26+ (Metal 4 — the MPP header live-compile requires it). /// Cached once; explicit FFAI_MOE_BGEMM_MPP=1 bypasses this. - private static let mppAutoCapable: Bool = { + /// + /// Module-visible (not `private`) rather than duplicated: `MoELayer`'s + /// F-85 coop-core MoE gather gate (`coopEnabled`) also wraps a + /// `matmul2d`/cooperative-tensor kernel (`ffai_moe_gather_qmm_coop`) + /// with the same Family-9/macOS-26 requirement, and unlike + /// `dispatchQmmMma`'s coop tier below (which falls back to a cheaper + /// tier on incapable hardware), that kernel has no lower-level + /// fallback of its own once dispatched - the capability check has to + /// happen at the `MoELayer` call site, so it reuses this same cached + /// check instead of re-deriving it. + static let mppAutoCapable: Bool = { guard ProcessInfo.processInfo.isOperatingSystemAtLeast( OperatingSystemVersion(majorVersion: 26, minorVersion: 0, patchVersion: 0)) @@ -6906,7 +6924,7 @@ public enum Ops { /// Multi-token Gated Delta Net recurrence over a chunk of `T` tokens. /// - /// Wraps `ffai_gated_delta_chunk` — same recurrence math as + /// Wraps `ffai_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 @@ -6919,17 +6937,17 @@ public enum Ops { /// * `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 + /// * `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 + /// 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 `ffai_gated_delta_step` apart from the + /// per group - identical to `ffai_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( @@ -8456,6 +8474,18 @@ public enum Ops { encoder.dispatchThreadgroups(tgGrid, threadsPerThreadgroup: tg) } + /// `FFAI_QMM_MPP` / `FFAI_QMM_COOP` raw env strings, cached once. + /// `dispatchQmmMma` below is the shared entry for every dense int4 + /// projection at prefill (GDN in-projections, attention QKV/out, + /// every other quantized linear), so this fires many times per layer + /// per forward call; re-fetching the whole process environment + /// dictionary that often was pure overhead once the tier decision + /// stabilizes for the life of the process. + private static let qmmMppEnvRaw: String? = + ProcessInfo.processInfo.environment["FFAI_QMM_MPP"] + private static let qmmCoopEnvRaw: String? = + ProcessInfo.processInfo.environment["FFAI_QMM_COOP"] + /// Inner dispatcher for `dequantGemmDynamicM`. Grid is `[N/32, M/32, 1]` /// × tg `[128, 1, 1]` (4 SGs WM=WN=2 per the canonical `ffai_qmm_mma`). /// `dispatchThreads` counts total threads per axis, so grid.x = N/32·128. @@ -8515,16 +8545,15 @@ public enum Ops { let kU = UInt32(k) let nU = UInt32(n) let gsU = UInt32(gsPerRow) - let env = ProcessInfo.processInfo.environment let mppRequested4: Bool - switch env["FFAI_QMM_MPP"] { + switch Self.qmmMppEnvRaw { case "0": mppRequested4 = false case "1": mppRequested4 = true default: mppRequested4 = Self.mppAutoCapable } let useMpp4 = mppRequested4 && gsPerRow > 0 && k % gsPerRow == 0 && k / gsPerRow == 64 let coopRequested4: Bool - switch env["FFAI_QMM_COOP"] { + switch Self.qmmCoopEnvRaw { case "0": coopRequested4 = false case "1": coopRequested4 = true default: coopRequested4 = Self.mppAutoCapable