Initial commit: E4B-MarkBase model integration with passing tests
CI / build-and-test (push) Has been cancelled
CI / build-and-test (push) Has been cancelled
- E4B-MarkBase model (42 layers, 4.4GB) loaded successfully - All Phase 1-6 tests passed (model loading, forward pass, vision/audio towers, token generation, performance) - All stress tests passed (5/5 in 127.6s) - Concurrent inference - Memory stress (67.5 tok/s, 0 NaN) - Continuous generation - Batch processing - Long-running stability - Swift Metal inference engine with multimodal support
This commit is contained in:
@@ -0,0 +1,188 @@
|
||||
import XCTest
|
||||
@testable import MarkBase
|
||||
|
||||
final class MoEExpertComputationTest: XCTestCase {
|
||||
|
||||
func testSingleExpertFusedGateUp() throws {
|
||||
print("\n═══════════════════════════════════════")
|
||||
print(" Single Expert Fused Gate+Up Test")
|
||||
print("═══════════════════════════════════════\n")
|
||||
|
||||
let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
|
||||
|
||||
print("Step 1: Load model...")
|
||||
let engine = try MarkBaseEngine(autoCompile: true)
|
||||
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 32)
|
||||
print(" ✓ Model loaded")
|
||||
|
||||
print("\nStep 2: Get expert components from layer 0...")
|
||||
let layer0 = model.layers[0]
|
||||
|
||||
guard let eGate = layer0.expertGate,
|
||||
let eUp = layer0.expertUp,
|
||||
let eDown = layer0.expertDown else {
|
||||
XCTFail("Expert components not present")
|
||||
return
|
||||
}
|
||||
|
||||
print(" ✓ Expert components retrieved")
|
||||
print(" Gate: \(eGate.numExperts) experts, outDim=\(eGate.expertOutDim)")
|
||||
print(" Up: \(eUp.numExperts) experts, outDim=\(eUp.expertOutDim)")
|
||||
print(" Down: \(eDown.numExperts) experts, outDim=\(eDown.expertOutDim)")
|
||||
|
||||
print("\nStep 3: Create minimal buffers...")
|
||||
let hs = model.hiddenSize
|
||||
let moeIntermediate = eGate.expertOutDim // 704
|
||||
|
||||
// Input buffer (normalized)
|
||||
let inputBuffer = engine.device.makeBuffer(length: hs * MemoryLayout<Float>.size)!
|
||||
var inputData = [Float](repeating: 0.1, count: hs)
|
||||
memcpy(inputBuffer.contents(), &inputData, inputData.count * MemoryLayout<Float>.size)
|
||||
|
||||
// Output buffer for gate+up (2 * moeIntermediate)
|
||||
let gateUpOutputBuffer = engine.device.makeBuffer(length: 2 * moeIntermediate * MemoryLayout<Float>.size)!
|
||||
|
||||
print(" ✓ Buffers created")
|
||||
print(" Input: \(hs) floats")
|
||||
print(" Gate+Up output: \(2 * moeIntermediate) floats")
|
||||
|
||||
print("\nStep 4: Test single expert (expert 0) gate+up...")
|
||||
let start = Date()
|
||||
|
||||
do {
|
||||
print(" Calling expertFusedGateUp for expert 0...")
|
||||
fflush(stdout)
|
||||
|
||||
let cmdBuf = engine.commandQueue.makeCommandBuffer()!
|
||||
try layer0.expertFusedGateUp(
|
||||
engine: engine,
|
||||
cmdBuf: cmdBuf,
|
||||
input: inputBuffer,
|
||||
gate: eGate,
|
||||
up: eUp,
|
||||
expertIdx: 0, // First expert only
|
||||
output: gateUpOutputBuffer
|
||||
)
|
||||
|
||||
print(" Command buffer created, committing...")
|
||||
fflush(stdout)
|
||||
|
||||
cmdBuf.commit()
|
||||
print(" Waiting for completion...")
|
||||
fflush(stdout)
|
||||
|
||||
cmdBuf.waitUntilCompleted()
|
||||
|
||||
let elapsed = Date().timeIntervalSince(start)
|
||||
print(" ✓ Expert gate+up completed in \(String(format: "%.3f", elapsed))s")
|
||||
print(" Command buffer status: \(cmdBuf.status.rawValue)")
|
||||
|
||||
if elapsed > 10.0 {
|
||||
print(" ⚠️ Expert took >10s - potential issue")
|
||||
}
|
||||
|
||||
if cmdBuf.status != .completed {
|
||||
let errorMsg = cmdBuf.error?.localizedDescription ?? "unknown"
|
||||
print(" ❌ Expert command failed: \(errorMsg)")
|
||||
XCTFail("Expert computation failed")
|
||||
return
|
||||
}
|
||||
|
||||
// Read output
|
||||
let gateUpOutput = engine.readFloats(from: gateUpOutputBuffer, count: 10)
|
||||
print("\n Gate+Up output (first 10): \(gateUpOutput)")
|
||||
|
||||
let hasNaN = gateUpOutput.contains { $0.isNaN }
|
||||
if hasNaN {
|
||||
print(" ❌ NaN detected in expert output!")
|
||||
XCTFail("Expert output has NaN")
|
||||
} else {
|
||||
print(" ✓ Expert output valid (no NaN)")
|
||||
}
|
||||
|
||||
} catch {
|
||||
print(" ❌ Expert computation failed: \(error)")
|
||||
throw error
|
||||
}
|
||||
|
||||
print("\n═══════════════════════════════════════")
|
||||
print("✓ Single expert gate+up works!")
|
||||
print("═══════════════════════════════════════\n")
|
||||
}
|
||||
|
||||
func testExpertLoopIteration() throws {
|
||||
print("\n═══════════════════════════════════════")
|
||||
print(" Expert Loop Iteration Test")
|
||||
print("═══════════════════════════════════════\n")
|
||||
|
||||
let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
|
||||
|
||||
print("Loading model...")
|
||||
let engine = try MarkBaseEngine(autoCompile: true)
|
||||
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 32)
|
||||
|
||||
let layer0 = model.layers[0]
|
||||
guard let eGate = layer0.expertGate, let eUp = layer0.expertUp else {
|
||||
XCTFail("Experts not present")
|
||||
return
|
||||
}
|
||||
|
||||
let hs = model.hiddenSize
|
||||
let moeIntermediate = eGate.expertOutDim
|
||||
|
||||
let inputBuffer = engine.device.makeBuffer(length: hs * MemoryLayout<Float>.size)!
|
||||
var inputData = [Float](repeating: 0.1, count: hs)
|
||||
memcpy(inputBuffer.contents(), &inputData, inputData.count * MemoryLayout<Float>.size)
|
||||
|
||||
let gateUpOutputBuffer = engine.device.makeBuffer(length: 2 * moeIntermediate * MemoryLayout<Float>.size)!
|
||||
|
||||
print("\nTesting multiple expert iterations...")
|
||||
|
||||
// Test iterating through 8 experts (top-k for 26B-A4B)
|
||||
let topKExperts = [0, 1, 2, 3, 4, 5, 6, 7]
|
||||
|
||||
for (i, expertIdx) in topKExperts.enumerated() {
|
||||
print(" Expert \(i): Calling expertFusedGateUp for expert \(expertIdx)...")
|
||||
fflush(stdout)
|
||||
|
||||
let start = Date()
|
||||
let cmdBuf = engine.commandQueue.makeCommandBuffer()!
|
||||
|
||||
do {
|
||||
try layer0.expertFusedGateUp(
|
||||
engine: engine,
|
||||
cmdBuf: cmdBuf,
|
||||
input: inputBuffer,
|
||||
gate: eGate,
|
||||
up: eUp,
|
||||
expertIdx: expertIdx,
|
||||
output: gateUpOutputBuffer
|
||||
)
|
||||
|
||||
cmdBuf.commit()
|
||||
cmdBuf.waitUntilCompleted()
|
||||
|
||||
let elapsed = Date().timeIntervalSince(start)
|
||||
print(" ✓ Expert \(expertIdx) completed in \(String(format: "%.3f", elapsed))s")
|
||||
|
||||
if elapsed > 5.0 {
|
||||
print(" ⚠️ Expert \(expertIdx) slow (>5s)")
|
||||
}
|
||||
|
||||
if cmdBuf.status != .completed {
|
||||
print(" ❌ Expert \(expertIdx) failed!")
|
||||
XCTFail("Expert \(expertIdx) iteration failed")
|
||||
return
|
||||
}
|
||||
|
||||
} catch {
|
||||
print(" ❌ Expert \(expertIdx) error: \(error)")
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
print("\n═══════════════════════════════════════")
|
||||
print("✓ Expert loop iteration works (8 experts tested)")
|
||||
print("═══════════════════════════════════════\n")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user