perf(engine/metal): gelu-fold wiring — token-identical, routed OFF by receipt (#341 phase 1)

The fused down projections (lthn_gelu_qmv local + lthn_gather_qmv_gelu
expert) wire in behind geluFoldEnabled with an engagement counter and a
session-level oracle: TestLoadGemma4QuantMoEGeluFoldMatchesChain proves
the fold token-identical to the gelu-dispatch + plain-down chain on the
QAT-shaped fixture (8-bit local + 4-bit expert downs both engaged) — the
verbatim load replication held, no FMA-contraction drift.

MEASURED AND ROUTED OFF: the chain evaluates each gelu ONCE; an x-load
fold re-evaluates it per row-tile — 352 threadgroups x 2112
precise::tanh on the 26B local down and x8 routes on the expert gather
(~6M redundant tanh per layer) — and the real 26B-A4B decode paid 136.4
tok/s fused vs 154.7 chained. geluFoldEnabled defaults false with the
receipt in-comment. LESSON (with the router-fusion receipt, now a rule):
fold only O(output) work into matvec loads; per-element work consumed by
many tiles must stay its own dispatch.

Receipts: parity oracle green (opt-in), 26B restored 155.0 tok/s with
the lane off, metal 1396 green, repo 10450 green.

Co-Authored-By: Virgil <virgil@lethean.io>

Author: Snider

go/engine/metal/dispatch_sink.go

@@ -440,6 +440,25 @@ func emitQMV[S dispatchSink](sink S, pso metal.MTLComputePipelineState, wq metal
 	emitQMVAt(sink, pso, wq, wqOff, scales, scalesOff, biases, biasesOff, x, 0, out, outOff, inDim, outDim)
 }
 
+// emitGeluQMV records the gelu-fused down projection (out = dequant(W) ·
+// gelu(gate)·up, lthn_gelu_qmv kernel — #341 phase 1) through any sink: wq=0,
+// scales=1, biases=2, gate=3, up=4, out=5, K=6, N=7, the same grid as emitQMV.
+// The separate gelu-gate-mul dispatch and its dependency hop never encode.
+func emitGeluQMV[S dispatchSink](sink S, pso metal.MTLComputePipelineState, wq metal.MTLBuffer, wqOff uint, scales metal.MTLBuffer, scalesOff uint, biases metal.MTLBuffer, biasesOff uint, gate, up, out metal.MTLBuffer, outOff uint, inDim, outDim int) {
+	sink.setPSO(pso)
+	sink.setBuf(wq, wqOff, 0)
+	sink.setBuf(scales, scalesOff, 1)
+	sink.setBuf(biases, biasesOff, 2)
+	sink.setBuf(gate, 0, 3)
+	sink.setBuf(up, 0, 4)
+	sink.setBuf(out, outOff, 5)
+	sink.setI32(int32(inDim), 6)  // K
+	sink.setI32(int32(outDim), 7) // N
+	const bn, bk = 8, 32
+	nTgp := uint((outDim + bn - 1) / bn)
+	sink.dispatchThreadgroups(metal.MTLSize{Width: 1, Height: nTgp, Depth: 1}, metal.MTLSize{Width: bk, Height: 2, Depth: 1})
+}
+
 // emitQMVAt is emitQMV with the activation vector bound at a byte offset — the batched dense
 // forward's rows live at byte offsets inside shared K-row buffers, and a row's quant gate reads
 // its input in place.

go/engine/metal/moe_block.go

@@ -1555,6 +1555,50 @@ func moeArriveZeroBuffer() metal.MTLBuffer {
 // the lane (a silent gate regression reads as zero, not as a perf blur).
 var moeConcurrentBlocks atomic.Int64
 
+// geluFoldDispatches counts MoE blocks that ran the gelu-fused down kernels
+// (#341 phase 1) instead of the gelu-dispatch + plain-down chain. Engagement
+// counter for the A/B tests.
+var geluFoldDispatches atomic.Int64
+
+type lthnGeluQMVKey struct {
+	groupSize, bits int
+}
+
+var (
+	lthnGeluQMVPSOMu    sync.Mutex
+	lthnGeluQMVPSOCache = map[lthnGeluQMVKey]metal.MTLComputePipelineState{}
+)
+
+// lthnGeluQMVPipeline resolves (and caches, including failures) the local-down
+// variant with the MLP gate fused into its x-load (lthn_gelu_qmv, #341 phase
+// 1). A miss — custom library absent, or a gs/bits pair outside the
+// instantiated set — caches nil so the block falls back to the gelu-dispatch +
+// plain-qmv chain without re-probing.
+func lthnGeluQMVPipeline(groupSize, bits int) (metal.MTLComputePipelineState, bool) {
+	key := lthnGeluQMVKey{groupSize: groupSize, bits: bits}
+	lthnGeluQMVPSOMu.Lock()
+	defer lthnGeluQMVPSOMu.Unlock()
+	if pso, ok := lthnGeluQMVPSOCache[key]; ok {
+		return pso, pso != nil
+	}
+	if customLibrary == nil || customLibrary.GetID() == 0 {
+		lthnGeluQMVPSOCache[key] = nil
+		return nil, false
+	}
+	fn := customLibrary.NewFunctionWithName(core.Sprintf("lthn_gelu_qmv_bfloat16_t_gs_%d_b_%d", groupSize, bits))
+	if fn == nil || fn.GetID() == 0 {
+		lthnGeluQMVPSOCache[key] = nil
+		return nil, false
+	}
+	pso, perr := device.NewComputePipelineStateWithFunctionError(fn)
+	if perr != nil {
+		lthnGeluQMVPSOCache[key] = nil
+		return nil, false
+	}
+	lthnGeluQMVPSOCache[key] = pso
+	return pso, true
+}
+
 // ffnMegaKernelCompatible is the KERNEL truth: the widths the megakernel is specialised for
 // (4/8-bit byte-aligned codes via ffnMegaPipelineBits, parity-proven at both) with all three
 // projections agreeing so one PSO serves the dispatch.
@@ -1706,6 +1750,20 @@ func encMoEBlockQuantDevice(enc metal.MTLComputeCommandEncoderObject, cb metal.M
 	if err != nil {
 		return enc, encConc, false, nil
 	}
+	// gelu fold (#341 phase 1): both down projections read gate/up directly and
+	// compute gelu(gate)·up at load — the two gelu dispatches (and the expert
+	// gelu's barrier) never encode. Decided ONCE for the whole block so the
+	// gated scratch stays coherent; either PSO missing (stale metallib, exotic
+	// width) or the lever falls back to the chain unchanged.
+	geluFold := false
+	var localDownGeluPSO, gatherDownGeluPSO metal.MTLComputePipelineState
+	if geluFoldEnabled {
+		if p1, ok1 := lthnGeluQMVPipeline(localDownGroupSize, localDownBits); ok1 {
+			if p2, ok2 := lthnGatherQMVPipeline(lthnGatherQMVKey{groupSize: expDownGroupSize, bits: expDownBits, expertRows: dModel, batchedX: true, gelu: true}); ok2 {
+				localDownGeluPSO, gatherDownGeluPSO, geluFold = p1, p2, true
+			}
+		}
+	}
 	rmsPSO, err := pipelineFor(rmsKernelBF16(dModel))
 	if err != nil {
 		return enc, encConc, false, nil
@@ -1881,8 +1939,22 @@ func encMoEBlockQuantDevice(enc metal.MTLComputeCommandEncoderObject, cb metal.M
 		}
 	}
 	emitGatherDownAll := func() {
+		if geluFold {
+			// the fused down reads each route's gate/up rows and gelus at load —
+			// the expert gelu dispatch never encoded.
+			emitLthnGatherQMVGeluRoutes(sink, gatherDownGeluPSO, scratch.expertGateAll, 0, scratch.expertUpAll, 0, expDownPacked.buf, expDownPacked.off, expDownScales.buf, expDownScales.off, expDownBiases.buf, expDownBiases.off, scratch.routeIota, routeIdxBuf, 0, scratch.expertDownAll, 0, dModel, expertDFF, expDownGroupSize, expDownBits, 0, topK)
+			return
+		}
 		emitGatherQMVAllRoutes(sink, gatherExpertDownPSO, downAllMeta, downKeyBatched, scratch.expertGatedAll, 0, expDownPacked.buf, expDownPacked.off, expDownScales.buf, expDownScales.off, expDownBiases.buf, expDownBiases.off, scratch.routeIota, routeIdxBuf, 0, scratch.expertDownAll, 0, dModel, expertDFF, expDownGroupSize, expDownBits, 0, topK)
 	}
+	emitLocalDown := func() {
+		if geluFold {
+			geluFoldDispatches.Add(1)
+			emitGeluQMV(sink, localDownGeluPSO, localDownPacked.buf, localDownPacked.off, localDownScales.buf, localDownScales.off, localDownBiases.buf, localDownBiases.off, msc.gate, msc.up, scratch.localOut, 0, dFF, dModel)
+			return
+		}
+		emitQ(localDownPSO, localDownPacked, localDownScales, localDownBiases, msc.gated, scratch.localOut, dFF, dModel)
+	}
 
 	// ---- concurrent pass: dispatches overlap; barriers mark the true dependency edges.
 	// The router, local-MLP, and expert branches are independent until the combine, so the
@@ -1921,19 +1993,27 @@ func encMoEBlockQuantDevice(enc metal.MTLComputeCommandEncoderObject, cb metal.M
 		emitQ(localGatePSO, localGatePacked, localGateScales, localGateBiases, scratch.localIn, msc.gate, dModel, dFF)
 		emitQ(localUpPSO, localUpPacked, localUpScales, localUpBiases, scratch.localIn, msc.up, dModel, dFF)
 		barrier()
-		// stage 3: top-k select ∥ local gelu
+		// stage 3: top-k select ∥ local gelu (folded into the local down when the
+		// fused kernels are available — the dispatch never encodes)
 		if !routerFused {
 			routerPlan.emitTopK(sink)
 		}
-		emitBinary(sink, geluPSO, msc.gate, 0, msc.up, 0, msc.gated, 0, dFF)
+		if !geluFold {
+			emitBinary(sink, geluPSO, msc.gate, 0, msc.up, 0, msc.gated, 0, dFF)
+		}
 		barrier()
 		// stage 4: expert gate/up gathers (need the top-k indices) ∥ local down
+		// (gelu-fused: reads msc.gate/msc.up from stage 2, two barriers upstream)
 		emitGatherInAll()
-		emitQ(localDownPSO, localDownPacked, localDownScales, localDownBiases, msc.gated, scratch.localOut, dFF, dModel)
-		barrier()
-		// stage 5: expert gelu
-		emitBinary(sink, geluPSO, scratch.expertGateAll, 0, scratch.expertUpAll, 0, scratch.expertGatedAll, 0, topK*expertDFF)
+		emitLocalDown()
 		barrier()
+		// stage 5: expert gelu — folded into the down gather's load when fused
+		// (the stage and its barrier disappear; the stage-4 barrier already
+		// orders the gate/up slabs ahead of the down gather)
+		if !geluFold {
+			emitBinary(sink, geluPSO, scratch.expertGateAll, 0, scratch.expertUpAll, 0, scratch.expertGatedAll, 0, topK*expertDFF)
+			barrier()
+		}
 		// stage 6: expert down gather
 		emitGatherDownAll()
 		barrier()
@@ -1989,14 +2069,18 @@ func encMoEBlockQuantDevice(enc metal.MTLComputeCommandEncoderObject, cb metal.M
 	} else {
 		emitQ(localGatePSO, localGatePacked, localGateScales, localGateBiases, scratch.localIn, msc.gate, dModel, dFF)
 		emitQ(localUpPSO, localUpPacked, localUpScales, localUpBiases, scratch.localIn, msc.up, dModel, dFF)
-		emitBinary(sink, geluPSO, msc.gate, 0, msc.up, 0, msc.gated, 0, dFF)
-		emitQ(localDownPSO, localDownPacked, localDownScales, localDownBiases, msc.gated, scratch.localOut, dFF, dModel)
+		if !geluFold {
+			emitBinary(sink, geluPSO, msc.gate, 0, msc.up, 0, msc.gated, 0, dFF)
+		}
+		emitLocalDown()
 	}
 	seam("moe.expert")
 	emitRMS(hBuf, pre2Buf.buf, scratch.expertIn, pre2Buf.off)
 	if allRoutes {
 		emitGatherInAll()
-		emitBinary(sink, geluPSO, scratch.expertGateAll, 0, scratch.expertUpAll, 0, scratch.expertGatedAll, 0, topK*expertDFF)
+		if !geluFold {
+			emitBinary(sink, geluPSO, scratch.expertGateAll, 0, scratch.expertUpAll, 0, scratch.expertGatedAll, 0, topK*expertDFF)
+		}
 		emitGatherDownAll()
 		seam("moe.tail")
 		// combine the routes: acc = Σ_r w_r · down_r — one fused dispatch (byte-identical

go/engine/metal/moe_session_test.go

@@ -815,3 +815,69 @@ func alignedSafetensorsBlob(t *testing.T, tensors map[string]safetensors.Tensor,
 	t.Fatal("could not build an aligned safetensors fixture")
 	return nil
 }
+
+// TestLoadGemma4QuantMoEGeluFoldMatchesChain pins the gelu-fused down
+// projections (#341 phase 1) to the gelu-dispatch + plain-down chain token for
+// token: the fused kernels compute T(gelu(gate)·up) at load with the chain's
+// exact expression and rounding (lthn_gelu_qmv_impl.h), so the fold is a
+// schedule-only change. The fixture's QAT-shaped overrides put the local down
+// at 8-bit and the expert down at 4-bit — both fused widths engage.
+func TestLoadGemma4QuantMoEGeluFoldMatchesChain(t *testing.T) {
+	if os.Getenv(MetallibPathEnv) == "" {
+		t.Skip("metallib not set")
+	}
+	const dModel, nHeads, nKV, headDim, vocab = 64, 2, 1, 64, 32
+	const dFF, expertDFF, numExperts, topK, numLayers = 128, 64, 4, 2, 2
+	const maxLen, n = 16, 6
+	quant := &model.QuantConfig{GroupSize: 64, Bits: 4, Overrides: map[string]model.ModuleQuant{}}
+	for i := range numLayers {
+		for _, m := range []string{"mlp.gate_proj", "mlp.up_proj", "mlp.down_proj", "router.proj"} {
+			quant.Overrides[core.Sprintf("model.layers.%d.%s", i, m)] = model.ModuleQuant{GroupSize: 64, Bits: 8}
+		}
+	}
+	cfg := g4.Config{
+		HiddenSize: dModel, NumHiddenLayers: numLayers, IntermediateSize: dFF,
+		NumAttentionHeads: nHeads, NumKeyValueHeads: nKV, HeadDim: headDim, VocabSize: vocab, RMSNormEps: 1e-6,
+		EnableMoEBlock: true, NumExperts: numExperts, TopKExperts: topK, MoEIntermediateSize: expertDFF,
+		Quantization: quant,
+	}
+	arch, err := cfg.Arch()
+	if err != nil {
+		t.Fatalf("Arch: %v", err)
+	}
+	ts := moeQuantTensors(t, arch, quant)
+	prompt := []int32{1, 5, 3}
+
+	gen := func(fold bool) []int32 {
+		t.Helper()
+		geluFoldEnabled = fold
+		defer func() { geluFoldEnabled = false }()
+		lm, aerr := model.Assemble(ts, arch, model.StandardWeightNames())
+		if aerr != nil {
+			t.Fatalf("model.Assemble: %v", aerr)
+		}
+		g, qerr := loadedToQuant(lm, quant.GroupSize, quant.Bits)
+		if qerr != nil {
+			t.Fatalf("loadedToQuant: %v", qerr)
+		}
+		sess, serr := NewArchQuantSession(g, arch, maxLen)
+		if serr != nil {
+			t.Fatalf("NewArchQuantSession: %v", serr)
+		}
+		out, gerr := sess.Generate(prompt, n, -1)
+		if gerr != nil {
+			t.Fatalf("Generate: %v", gerr)
+		}
+		return out
+	}
+	chain := gen(false)
+	before := geluFoldDispatches.Load()
+	folded := gen(true)
+	if geluFoldDispatches.Load() == before {
+		t.Skip("gelu-fused down kernels unavailable (stale metallib) — the compare would be vacuous")
+	}
+	if !idsEqual(folded, chain) {
+		t.Fatalf("gelu fold diverged from the chain: %v != %v", folded, chain)
+	}
+	t.Logf("gelu-fused downs ≡ chain over %d tokens (%v)", len(chain), chain)
+}

go/engine/metal/piece_timing.go

@@ -48,6 +48,16 @@ var (
 	// lthn_gather_qmv kernel is available (7 buffers + 2 scalars, #280). Test/bench hook
 	// for lean-vs-MLX byte compares.
 	leanGatherDisabled bool
+	// geluFoldEnabled opts the MoE down projections into the gelu-fused kernels
+	// (gelu(gate)·up computed at the down's x-load — token-identical to the
+	// chain, see lthn_gelu_qmv_impl.h). OFF by receipt: the chain evaluates each
+	// gelu ONCE, but an x-load fold re-evaluates it per row-tile — 352
+	// threadgroups × 2112 precise::tanh for the 26B local down and ×8 routes
+	// for the expert gather (~6M redundant tanh per layer) — and the real
+	// 26B-A4B decode measured 136.4 tok/s fused vs 154.7 chained. Fold only
+	// O(output) work into matvec loads; per-element work consumed by many tiles
+	// must stay its own dispatch. Capability + parity oracle kept.
+	geluFoldEnabled bool
 	// encCarryDisabled forces every concurrent pass to close its encoder and
 	// reopen a serial one on exit (the pre-#341 shape) instead of carrying the
 	// open concurrent encoder into the next pass. Test/bench hook for

go/engine/metal/qmv_gather.go

@@ -185,6 +185,10 @@ var leanGatherDispatches atomic.Int64
 type lthnGatherQMVKey struct {
 	groupSize, bits, expertRows int
 	fast, batchedX              bool
+	// gelu selects the down-projection variant with the MLP gate fused into its
+	// x-load (lthn_gather_qmv_gelu — #341 phase 1). Only the qmv_impl form is
+	// instantiated, so gelu keys always carry fast=false.
+	gelu bool
 }
 
 var (
@@ -211,6 +215,9 @@ func lthnGatherQMVPipeline(key lthnGatherQMVKey) (metal.MTLComputePipelineState,
 	if key.fast {
 		variant = "lthn_gather_qmv_fast_"
 	}
+	if key.gelu {
+		variant = "lthn_gather_qmv_gelu_"
+	}
 	name := core.Sprintf("%sbfloat16_t_gs_%d_b_%d", variant, key.groupSize, key.bits)
 	fc := metal.NewMTLFunctionConstantValues()
 	rows := int32(key.expertRows)
@@ -255,6 +262,31 @@ func emitLthnGatherQMVRoutes(sink encSink, pso metal.MTLComputePipelineState, x
 	)
 }
 
+// emitLthnGatherQMVGeluRoutes is emitLthnGatherQMVRoutes for the gelu-fused
+// down variant (#341 phase 1): the kernel reads each route's gate/up activation
+// rows (same lhs indexing on both) and computes gelu(gate)·up at load, so the
+// expert gelu dispatch and its barrier never encode.
+func emitLthnGatherQMVGeluRoutes(sink encSink, pso metal.MTLComputePipelineState, gate metal.MTLBuffer, gateOff uint, up metal.MTLBuffer, upOff uint, wq metal.MTLBuffer, wqOff uint, scales metal.MTLBuffer, scalesOff uint, biases metal.MTLBuffer, biasesOff uint, lhsIdx, rhsIdx metal.MTLBuffer, rhsOff uint, out metal.MTLBuffer, outOff uint, outDim, inDim, groupSize, bits, rowBase, routes int) {
+	rowPackedBytes := inDim * bits / 8
+	groups := inDim / groupSize
+	sink.setPSO(pso)
+	sink.setBuf(wq, wqOff+uint(rowBase*rowPackedBytes), 0)
+	sink.setBuf(scales, scalesOff+uint(rowBase*groups*bf16Size), 1)
+	sink.setBuf(biases, biasesOff+uint(rowBase*groups*bf16Size), 2)
+	sink.setBuf(gate, gateOff, 3)
+	sink.setBuf(up, upOff, 4)
+	sink.setBuf(lhsIdx, 0, 5)
+	sink.setBuf(rhsIdx, rhsOff, 6)
+	sink.setBuf(out, outOff, 7)
+	sink.setI32(int32(inDim), 8)
+	sink.setI32(int32(outDim), 9)
+	const bn, bk = 8, 32
+	sink.dispatchThreadgroups(
+		metal.MTLSize{Width: 1, Height: uint((outDim + bn - 1) / bn), Depth: uint(routes)},
+		metal.MTLSize{Width: bk, Height: 2, Depth: 1},
+	)
+}
+
 // emitGatherQMVAllRoutes dispatches one gather_qmv covering EVERY route: grid.z carries the
 // routes, rhs_indices supplies each route's expert id (the router's device idxBuf), and the
 // x/lhs binding either shares one row (gate/up) or walks the per-route gated rows (down).