Add bf16 layer weight support for E4B model

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