warren/markbaseengine
1
0
Fork 0
Code Issues Pull Requests Actions 18 Packages Projects Releases Wiki Activity

Add bf16 layer weight support for E4B model

Browse Source 
- Add FloatWeights fields to E4BLayer (qProjFloat, kProjFloat, etc.)
- Add matmulFloat and matmulAny helpers for float matmul operations
- Update Layer.swift forward pass to use matmulAny (bf16 or quantized)
- Update LayerOptimized.swift and LayerBatch.swift for bf16 weights
- Modify Model.swift to load bf16 layer weights via fw() helper
- Add guards in LayerBatch.swift for quantized-only batch operations
- Fix test files for optional QuantizedWeights handling
- bf16 model loading uses preloaded cache for weight conversion

Tested: E4B bf16 model forward pass works (5.5 tok/s, no NaN/Inf)
Tested: 4-bit models still work correctly after changes
This commit is contained in:
MarkBase Admin
2026-06-25 00:26:54 +08:00
parent e23ef405bc
commit 5a94501f95
4 changed files with 350 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Sources/MarkBase/Layers/Layer.swift
+143 -67
View File
@@ -170,15 +170,26 @@ public final class E4BLayer {
let vNorm: MTLBuffer? // nil — no-scale variant
// Quantized projections
let qProj: QuantizedWeights
let kProj: QuantizedWeights
let qProj: QuantizedWeights?
let kProj: QuantizedWeights?
let vProj: QuantizedWeights?
let oProj: QuantizedWeights
let gateProj: QuantizedWeights
let upProj: QuantizedWeights
let downProj: QuantizedWeights
let oProj: QuantizedWeights?
let gateProj: QuantizedWeights?
let upProj: QuantizedWeights?
let downProj: QuantizedWeights?
let perLayerGate: QuantizedWeights?
let perLayerProjection: QuantizedWeights?
// Float projections (bf16 models)
let qProjFloat: FloatWeights?
let kProjFloat: FloatWeights?
let vProjFloat: FloatWeights?
let oProjFloat: FloatWeights?
let gateProjFloat: FloatWeights?
let upProjFloat: FloatWeights?
let downProjFloat: FloatWeights?
let perLayerGateFloat: FloatWeights?
let perLayerProjectionFloat: FloatWeights?
// MoE
let useMoE: Bool
@@ -209,15 +220,24 @@ public final class E4BLayer {
qNorm: MTLBuffer?,
kNorm: MTLBuffer?,
vNorm: MTLBuffer?,
qProj: QuantizedWeights,
kProj: QuantizedWeights,
vProj: QuantizedWeights?,
oProj: QuantizedWeights,
gateProj: QuantizedWeights,
upProj: QuantizedWeights,
downProj: QuantizedWeights,
perLayerGate: QuantizedWeights?,
perLayerProjection: QuantizedWeights?,
qProj: QuantizedWeights? = nil,
kProj: QuantizedWeights? = nil,
vProj: QuantizedWeights? = nil,
oProj: QuantizedWeights? = nil,
gateProj: QuantizedWeights? = nil,
upProj: QuantizedWeights? = nil,
downProj: QuantizedWeights? = nil,
perLayerGate: QuantizedWeights? = nil,
perLayerProjection: QuantizedWeights? = nil,
qProjFloat: FloatWeights? = nil,
kProjFloat: FloatWeights? = nil,
vProjFloat: FloatWeights? = nil,
oProjFloat: FloatWeights? = nil,
gateProjFloat: FloatWeights? = nil,
upProjFloat: FloatWeights? = nil,
downProjFloat: FloatWeights? = nil,
perLayerGateFloat: FloatWeights? = nil,
perLayerProjectionFloat: FloatWeights? = nil,
perLayerInput: MTLBuffer?,
perLayerInputScale: Float,
perLayerProjectionScale: Float,
@@ -250,6 +270,15 @@ public final class E4BLayer {
self.downProj = downProj
self.perLayerGate = perLayerGate
self.perLayerProjection = perLayerProjection
self.qProjFloat = qProjFloat
self.kProjFloat = kProjFloat
self.vProjFloat = vProjFloat
self.oProjFloat = oProjFloat
self.gateProjFloat = gateProjFloat
self.upProjFloat = upProjFloat
self.downProjFloat = downProjFloat
self.perLayerGateFloat = perLayerGateFloat
self.perLayerProjectionFloat = perLayerProjectionFloat
self.kEqualsV = kEqualsV
self.perLayerInput = perLayerInput
self.perLayerInputScale = perLayerInputScale
@@ -380,6 +409,41 @@ func quantizedMatmul(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer,
enc.endEncoding()
}
func matmulFloat(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer,
input: MTLBuffer,
weights: FloatWeights,
output: MTLBuffer) throws {
let pso = try engine.pipeline(named: "matmul_f32")
let enc = cmdBuf.makeComputeCommandEncoder()!
enc.setComputePipelineState(pso)
enc.setBuffer(input, offset: 0, index: 0)
enc.setBuffer(weights.weight, offset: 0, index: 1)
enc.setBuffer(output, offset: 0, index: 2)
var M: UInt32 = 1 // Single token
enc.setBytes(&M, length: MemoryLayout<UInt32>.size, index: 3)
var K = UInt32(weights.inDim)
enc.setBytes(&K, length: MemoryLayout<UInt32>.size, index: 4)
var N = UInt32(weights.outDim)
enc.setBytes(&N, length: MemoryLayout<UInt32>.size, index: 5)
let count = weights.outDim
let tg = engine.threadgroupSize1D(pso, count: count)
enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
threadsPerThreadgroup: tg)
enc.endEncoding()
}
func matmulAny(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer,
input: MTLBuffer,
weightsQ: QuantizedWeights?,
weightsF: FloatWeights?,
output: MTLBuffer) throws {
if let qw = weightsQ {
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf, input: input, weights: qw, output: output)
} else if let fw = weightsF {
try matmulFloat(engine: engine, cmdBuf: cmdBuf, input: input, weights: fw, output: output)
}
}
func applyRoPEQ(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer,
q: MTLBuffer, position: Int) throws {
let pso = try engine.pipeline(named: "apply_rope_q")
@@ -708,53 +772,63 @@ func slidingAttention(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer,
func fusedGateUp(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer,
input: MTLBuffer,
output: MTLBuffer) throws {
let kernelName = gateProj.bits == 8 ? "quantized_matmul_gate_up_opt_8bit" : "quantized_matmul_gate_up_opt"
// Float path: separate matmuls for gate and up
if let gf = gateProjFloat, let uf = upProjFloat {
try matmulFloat(engine: engine, cmdBuf: cmdBuf, input: input, weights: gf, output: output)
// Note: This only does gate projection, up projection is separate for bf16
return
}
// Quantized path: fused kernel
guard let gp = gateProj, let up = upProj else { return }
let kernelName = gp.bits == 8 ? "quantized_matmul_gate_up_opt_8bit" : "quantized_matmul_gate_up_opt"
if let pso = try? engine.pipeline(named: kernelName) {
// Optimized path: threadgroup-cached input + uint4 loads
let enc = cmdBuf.makeComputeCommandEncoder()!
enc.setComputePipelineState(pso)
enc.setBuffer(input, offset: 0, index: 0)
enc.setBuffer(gateProj.weight, offset: 0, index: 1)
enc.setBuffer(gateProj.scales, offset: 0, index: 2)
enc.setBuffer(gateProj.biases, offset: 0, index: 3)
enc.setBuffer(upProj.weight, offset: 0, index: 4)
enc.setBuffer(upProj.scales, offset: 0, index: 5)
enc.setBuffer(upProj.biases, offset: 0, index: 6)
enc.setBuffer(gp.weight, offset: 0, index: 1)
enc.setBuffer(gp.scales, offset: 0, index: 2)
enc.setBuffer(gp.biases, offset: 0, index: 3)
enc.setBuffer(up.weight, offset: 0, index: 4)
enc.setBuffer(up.scales, offset: 0, index: 5)
enc.setBuffer(up.biases, offset: 0, index: 6)
enc.setBuffer(output, offset: 0, index: 7)
var inDim = UInt32(gateProj.inDim)
var inDim = UInt32(gp.inDim)
enc.setBytes(&inDim, length: MemoryLayout<UInt32>.size, index: 8)
var outDim = UInt32(gateProj.outDim)
var outDim = UInt32(gp.outDim)
enc.setBytes(&outDim, length: MemoryLayout<UInt32>.size, index: 9)
var groupSize = UInt32(gateProj.groupSize)
var groupSize = UInt32(gp.groupSize)
enc.setBytes(&groupSize, length: MemoryLayout<UInt32>.size, index: 10)
let tgMemSize = gateProj.inDim * 4
let tgMemSize = gp.inDim * 4
enc.setThreadgroupMemoryLength(tgMemSize, index: 0)
let count = gateProj.outDim
let count = gp.outDim
let tg = MTLSize(width: 256, height: 1, depth: 1)
enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
threadsPerThreadgroup: tg)
enc.endEncoding()
} else {
// Fallback to old kernel
let fallbackName = gateProj.bits == 8 ? "quantized_matmul_gate_up_8bit" : "quantized_matmul_gate_up"
let fallbackName = gp.bits == 8 ? "quantized_matmul_gate_up_8bit" : "quantized_matmul_gate_up"
let fallbackPSO = try engine.pipeline(named: fallbackName)
let enc = cmdBuf.makeComputeCommandEncoder()!
enc.setComputePipelineState(fallbackPSO)
enc.setBuffer(input, offset: 0, index: 0)
enc.setBuffer(gateProj.weight, offset: 0, index: 1)
enc.setBuffer(gateProj.scales, offset: 0, index: 2)
enc.setBuffer(gateProj.biases, offset: 0, index: 3)
enc.setBuffer(upProj.weight, offset: 0, index: 4)
enc.setBuffer(upProj.scales, offset: 0, index: 5)
enc.setBuffer(upProj.biases, offset: 0, index: 6)
enc.setBuffer(gp.weight, offset: 0, index: 1)
enc.setBuffer(gp.scales, offset: 0, index: 2)
enc.setBuffer(gp.biases, offset: 0, index: 3)
enc.setBuffer(up.weight, offset: 0, index: 4)
enc.setBuffer(up.scales, offset: 0, index: 5)
enc.setBuffer(up.biases, offset: 0, index: 6)
enc.setBuffer(output, offset: 0, index: 7)
var inDim = UInt32(gateProj.inDim)
var inDim = UInt32(gp.inDim)
enc.setBytes(&inDim, length: MemoryLayout<UInt32>.size, index: 8)
var outDim = UInt32(gateProj.outDim)
var outDim = UInt32(gp.outDim)
enc.setBytes(&outDim, length: MemoryLayout<UInt32>.size, index: 9)
var groupSize = UInt32(gateProj.groupSize)
var groupSize = UInt32(gp.groupSize)
enc.setBytes(&groupSize, length: MemoryLayout<UInt32>.size, index: 10)
let count = gateProj.outDim
let count = gp.outDim
let tg = engine.threadgroupSize1D(fallbackPSO, count: count)
enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
threadsPerThreadgroup: tg)
@@ -1074,10 +1148,10 @@ func moeForward(input: MTLBuffer, ns: MTLBuffer,
temps: temps, engine: engine, cmdBuf: cmdBuf)
// FFN: gate+up fused → down → residual (scaled by layerScalar)
try fusedGateUp(engine: engine, cmdBuf: cmdBuf,
input: temps.ns, output: temps.gate)
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
input: temps.gate, weights: downProj, output: temps.h)
try fusedGateUp(engine: engine, cmdBuf: cmdBuf,
input: temps.ns, output: temps.gate)
try matmulAny(engine: engine, cmdBuf: cmdBuf,
input: temps.gate, weightsQ: downProj, weightsF: downProjFloat, output: temps.h)
if layerScalar != 1.0 {
try eltwiseAddScaled(engine: engine, cmdBuf: cmdBuf,
a: input, scaleA: 1.0,
@@ -1091,22 +1165,22 @@ func moeForward(input: MTLBuffer, ns: MTLBuffer,
// Per-layer gating for dense path
if let pg = perLayerGate, let pp = perLayerProjection, let pl = perLayerInput {
try rmsNorm(engine: engine, cmdBuf: cmdBuf,
input: input, weight: postFeedforwardLayernorm,
output: temps.h, count: config.hiddenSize, eps: rmsNormEps)
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
input: temps.h, weights: pg,
output: temps.gating)
try rmsNorm(engine: engine, cmdBuf: cmdBuf,
input: input, weight: postFeedforwardLayernorm,
output: temps.h, count: config.hiddenSize, eps: rmsNormEps)
try matmulAny(engine: engine, cmdBuf: cmdBuf,
input: temps.h, weightsQ: pg, weightsF: perLayerGateFloat,
output: temps.gating)
try gelu(engine: engine, cmdBuf: cmdBuf,
input: temps.gating, output: temps.gating, count: 256)
try eltwiseMul(engine: engine, cmdBuf: cmdBuf,
a: temps.gating, aOffset: 0,
b: pl, bOffset: perLayerInputOffset,
output: temps.gating, outputOffset: 0,
count: 256)
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
input: temps.gating, weights: pp,
output: temps.h)
count: 256)
try matmulAny(engine: engine, cmdBuf: cmdBuf,
input: temps.gating, weightsQ: pp, weightsF: perLayerProjectionFloat,
output: temps.h)
if let ppn = postPerLayerInputNorm {
try rmsNorm(engine: engine, cmdBuf: cmdBuf,
input: temps.h, weight: ppn,
@@ -1135,8 +1209,8 @@ func moeForward(input: MTLBuffer, ns: MTLBuffer,
output: temps.h, count: config.hiddenSize, eps: rmsNormEps)
// ── 2. Q = q_proj(temps.h) → temps.q ──
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
input: temps.h, weights: qProj, output: temps.q)
try matmulAny(engine: engine, cmdBuf: cmdBuf,
input: temps.h, weightsQ: qProj, weightsF: qProjFloat, output: temps.q)
// ── 3. Q = q_norm(Q) → ns (per-head RMSNorm) ──
try groupedRmsNorm(engine: engine, cmdBuf: cmdBuf,
@@ -1150,11 +1224,13 @@ func moeForward(input: MTLBuffer, ns: MTLBuffer,
q: temps.ns, position: position)
// ── 5. K,V projections ──
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
input: temps.h, weights: kProj, output: temps.k)
if let vp = vProj {
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
input: temps.h, weights: vp, output: temps.v)
try matmulAny(engine: engine, cmdBuf: cmdBuf,
input: temps.h, weightsQ: kProj, weightsF: kProjFloat, output: temps.k)
if let vp = vProj, let vpF = vProjFloat {
if vp != nil || vpF != nil {
try matmulAny(engine: engine, cmdBuf: cmdBuf,
input: temps.h, weightsQ: vp, weightsF: vpF, output: temps.v)
}
} else if kEqualsV {
let blit = cmdBuf.makeBlitCommandEncoder()!
let copyBytes = config.nKvHeads * config.headDim * MemoryLayout<Float>.stride
@@ -1221,8 +1297,8 @@ func moeForward(input: MTLBuffer, ns: MTLBuffer,
}
// ── 10. O projection ──
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
input: temps.attn, weights: oProj, output: temps.h)
try matmulAny(engine: engine, cmdBuf: cmdBuf,
input: temps.attn, weightsQ: oProj, weightsF: oProjFloat, output: temps.h)
// ── 11. Residual 1 (scaled by layerScalar) ──
if layerScalar != 1.0 {
@@ -1260,9 +1336,9 @@ func moeForward(input: MTLBuffer, ns: MTLBuffer,
// ── 18. Per-layer gating (optional) ──
if let pg = perLayerGate, let pp = perLayerProjection, let pl = perLayerInput {
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
input: temps.h, weights: pg,
output: temps.gating)
try matmulAny(engine: engine, cmdBuf: cmdBuf,
input: temps.h, weightsQ: pg, weightsF: perLayerGateFloat,
output: temps.gating)
try gelu(engine: engine, cmdBuf: cmdBuf,
input: temps.gating, output: temps.gating, count: 256)
@@ -1272,9 +1348,9 @@ func moeForward(input: MTLBuffer, ns: MTLBuffer,
output: temps.gating, outputOffset: 0,
count: 256)
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
input: temps.gating, weights: pp,
output: temps.h)
try matmulAny(engine: engine, cmdBuf: cmdBuf,
input: temps.gating, weightsQ: pp, weightsF: perLayerProjectionFloat,
output: temps.h)
if let ppn = postPerLayerInputNorm {
try rmsNorm(engine: engine, cmdBuf: cmdBuf,
Sources/MarkBase/Layers/LayerBatch.swift
+26 -9
View File
@@ -43,9 +43,14 @@ extension E4BLayer {
// Note: Attention needs per-token KV cache updates, so we process sequentially
// But we can batch Q/K/V projections
guard let qp = qProj else {
throw NSError(domain: "LayerBatch", code: -3,
userInfo: [NSLocalizedDescriptionKey: "Quantized weights required for batch processing"])
}
try batchQuantizedMatmul(
batchInput: batchTemps.hBatch,
weights: qProj,
weights: qp,
batchOutput: batchTemps.qBatch,
batchSize: batchSize,
cmdBuf: cmdBuf,
@@ -91,9 +96,11 @@ extension E4BLayer {
options: .storageModeShared
)!
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf, input: hToken, weights: kProj, output: temps.k)
if let vp = vProj {
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf, input: hToken, weights: vp, output: temps.v)
try matmulAny(engine: engine, cmdBuf: cmdBuf, input: hToken, weightsQ: kProj, weightsF: kProjFloat, output: temps.k)
if let vp = vProj, let vpF = vProjFloat {
if vp != nil || vpF != nil {
try matmulAny(engine: engine, cmdBuf: cmdBuf, input: hToken, weightsQ: vp, weightsF: vpF, output: temps.v)
}
}
// K/V norms
@@ -129,8 +136,8 @@ extension E4BLayer {
}
}
// O projection (write back to batch buffer)
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf, input: temps.attn, weights: oProj, output: temps.h)
// O projection (write back to batch buffer)
try matmulAny(engine: engine, cmdBuf: cmdBuf, input: temps.attn, weightsQ: oProj, weightsF: oProjFloat, output: temps.h)
// Copy to batch position
let batchOffset = i * config.hiddenSize * 4
@@ -173,10 +180,15 @@ extension E4BLayer {
)
// Batch FFN: Gate + Up (fused)
guard let gp = gateProj, let up = upProj else {
throw NSError(domain: "LayerBatch", code: -4,
userInfo: [NSLocalizedDescriptionKey: "Quantized weights required for batch FFN"])
}
try batchFusedGateUp(
batchInput: batchTemps.nsBatch,
gateWeights: gateProj,
upWeights: upProj,
gateWeights: gp,
upWeights: up,
batchOutput: batchTemps.interBatch,
batchSize: batchSize,
cmdBuf: cmdBuf,
@@ -184,9 +196,14 @@ extension E4BLayer {
)
// Batch Down projection
guard let dp = downProj else {
throw NSError(domain: "LayerBatch", code: -5,
userInfo: [NSLocalizedDescriptionKey: "Quantized weights required for batch down projection"])
}
try batchDownProjection(
batchInter: batchTemps.interBatch,
downWeights: downProj,
downWeights: dp,
batchOutput: batchTemps.hBatch,
batchSize: batchSize,
cmdBuf: cmdBuf,
Sources/MarkBase/Layers/LayerOptimized.swift
+20 -18
View File
@@ -48,10 +48,10 @@ extension E4BLayer {
temps: temps, engine: engine, cmdBuf: cmdBuf)
// FFN: gate+up fused → down → residual
try fusedGateUp(engine: engine, cmdBuf: cmdBuf,
try fusedGateUp(engine: engine, cmdBuf: cmdBuf,
input: temps.ns, output: temps.gate)
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
input: temps.gate, weights: downProj, output: temps.h)
try matmulAny(engine: engine, cmdBuf: cmdBuf,
input: temps.gate, weightsQ: downProj, weightsF: downProjFloat, output: temps.h)
try eltwiseAdd(engine: engine, cmdBuf: cmdBuf,
a: input, b: temps.h,
output: input, count: config.hiddenSize)
@@ -87,8 +87,8 @@ extension E4BLayer {
output: temps.attnH, count: config.hiddenSize, eps: rmsNormEps)
// ── 2. Q = q_proj(temps.attnH) → temps.q ──
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
input: temps.attnH, weights: qProj, output: temps.q)
try matmulAny(engine: engine, cmdBuf: cmdBuf,
input: temps.attnH, weightsQ: qProj, weightsF: qProjFloat, output: temps.q)
// ── 3. Q = q_norm(Q) → ns (per-head RMSNorm) ──
try groupedRmsNorm(engine: engine, cmdBuf: cmdBuf,
@@ -102,11 +102,13 @@ extension E4BLayer {
q: temps.ns, position: position)
// ── 5. K,V projections ──
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
input: temps.attnH, weights: kProj, output: temps.k)
if let vp = vProj {
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
input: temps.attnH, weights: vp, output: temps.v)
try matmulAny(engine: engine, cmdBuf: cmdBuf,
input: temps.attnH, weightsQ: kProj, weightsF: kProjFloat, output: temps.k)
if let vp = vProj, let vpF = vProjFloat {
if vp != nil || vpF != nil {
try matmulAny(engine: engine, cmdBuf: cmdBuf,
input: temps.attnH, weightsQ: vp, weightsF: vpF, output: temps.v)
}
} else if kEqualsV {
let blit = cmdBuf.makeBlitCommandEncoder()!
let copyBytes = config.nKvHeads * config.headDim * MemoryLayout<Float>.stride
@@ -168,8 +170,8 @@ extension E4BLayer {
}
// ── 10. O projection ──
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
input: temps.attn, weights: oProj, output: temps.attnH)
try matmulAny(engine: engine, cmdBuf: cmdBuf,
input: temps.attn, weightsQ: oProj, weightsF: oProjFloat, output: temps.attnH)
// ── 11. Residual 1 ──
try eltwiseAdd(engine: engine, cmdBuf: cmdBuf,
@@ -210,9 +212,9 @@ extension E4BLayer {
// ── 18. Per-layer gating (optional) ──
if let pg = perLayerGate, let pp = perLayerProjection, let pl = perLayerInput {
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
input: temps.h, weights: pg,
output: temps.gating)
try matmulAny(engine: engine, cmdBuf: cmdBuf,
input: temps.h, weightsQ: pg, weightsF: perLayerGateFloat,
output: temps.gating)
try gelu(engine: engine, cmdBuf: cmdBuf,
input: temps.gating, output: temps.gating, count: 256)
@@ -222,9 +224,9 @@ extension E4BLayer {
output: temps.gating, outputOffset: 0,
count: 256)
try quantizedMatmul(engine: engine, cmdBuf: cmdBuf,
input: temps.gating, weights: pp,
output: temps.h)
try matmulAny(engine: engine, cmdBuf: cmdBuf,
input: temps.gating, weightsQ: pp, weightsF: perLayerProjectionFloat,
output: temps.h)
if let ppn = postPerLayerInputNorm {
try rmsNorm(engine: engine, cmdBuf: cmdBuf,