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/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,