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

v2: fix 26B activation explosion — normalize groupSize=32 scales, fix hardcoded loops
CI / build (push) Waiting to run
Details
CI / unit-tests (push) Blocked by required conditions
Details
CI / lint (push) Blocked by required conditions
Details

Browse Source
This commit is contained in:
MarkBase Admin
2026-07-05 19:52:47 +08:00
parent 8a29dae613
commit 239474bef0
4 changed files with 60 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Sources/MarkBase/Layers/Layer.swift
+10 -6
View File
@@ -366,9 +366,8 @@ func quantizedMatmul(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer,
weights: QuantizedWeights,
output: MTLBuffer) throws {
// Select kernel based on quantization bits
let kernelName = weights.bits == 8 ? "quantized_matmul_8bit" : "quantized_matmul"
// TEMPORARILY USE FALLBACK KERNEL FOR TESTING
if false, let pso = try? engine.pipeline(named: kernelName) {
let kernelName = weights.bits == 8 ? "quantized_matmul_simd_8bit" : "quantized_matmul"
if let pso = try? engine.pipeline(named: kernelName) {
let enc = cmdBuf.makeComputeCommandEncoder()!
enc.setComputePipelineState(pso)
enc.setBuffer(input, offset: 0, index: 0)
@@ -868,7 +867,7 @@ func quantizedMatmulExpert(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer,
enc.setBytes(&inDim, length: MemoryLayout<UInt32>.size, index: 5)
var outDim = UInt32(expert.expertOutDim)
enc.setBytes(&outDim, length: MemoryLayout<UInt32>.size, index: 6)
var groupSize = UInt32(expert.expertInDim / 64)
var groupSize = UInt32(expert.expertInDim / expert.numGroups)
enc.setBytes(&groupSize, length: MemoryLayout<UInt32>.size, index: 7)
let tg = engine.threadgroupSize1D(fallbackPSO, count: expert.expertOutDim)
enc.dispatchThreads(MTLSize(width: expert.expertOutDim, height: 1, depth: 1),
@@ -922,7 +921,7 @@ func quantizedMatmulExpert(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer,
enc.setBytes(&inDim, length: MemoryLayout<UInt32>.size, index: 8)
var outDim = UInt32(gate.expertOutDim)
enc.setBytes(&outDim, length: MemoryLayout<UInt32>.size, index: 9)
var groupSize = UInt32(gate.expertInDim / 64) // group_size is 64 for quantized weights
var groupSize = UInt32(gate.expertInDim / gate.numGroups)
enc.setBytes(&groupSize, length: MemoryLayout<UInt32>.size, index: 10)
let count = gate.expertOutDim
let tg = engine.threadgroupSize1D(pso, count: count)
@@ -977,6 +976,10 @@ func quantizedMatmulExpert(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer,
gate: MoEExpertGroup, up: MoEExpertGroup, down: MoEExpertGroup,
accum: MTLBuffer) throws -> Bool {
guard let pso = try? engine.pipeline(named: "moe_mega_kernel") else { return false }
// Mega kernel supports only 4-bit router with groupSize=64 experts
guard router.bits == 4 else { return false }
let expertGroupSize = gate.expertInDim / gate.numGroups
guard expertGroupSize == 64 else { return false }
let enc = cmdBuf.makeComputeCommandEncoder()!
enc.setComputePipelineState(pso)
enc.setBuffer(input, offset: 0, index: 0)
@@ -1095,8 +1098,9 @@ func moeForward(input: MTLBuffer, ns: MTLBuffer,
expertIdx: expertIdx,
accum: temps.h, weight: weight)
}
}
// ── Step 5: Residual: input += moe_output (temps.h) scaled by layerScalar ──
if layerScalar != 1.0 {
try eltwiseAddScaled(engine: engine, cmdBuf: cmdBuf,