MarkBase Admin
ac75faa0cc
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
266 lines
9.4 KiB
Swift
266 lines
9.4 KiB
Swift
import XCTest
|
||
@testable import MarkBase
|
||
|
||
final class G12BOptimizationTests: XCTestCase {
|
||
|
||
// ── Buffer Pool Tests ─────────────────────────────
|
||
|
||
func testBufferPoolAcquireAndRelease() throws {
|
||
let engine = try MarkBaseEngine()
|
||
|
||
let buf1 = engine.acquireBuffer(length: 1024)
|
||
XCTAssertGreaterThan(buf1.length, 0)
|
||
|
||
engine.releaseBuffer(buf1)
|
||
|
||
let buf2 = engine.acquireBuffer(length: 1024)
|
||
XCTAssertEqual(buf1.length, buf2.length)
|
||
}
|
||
|
||
func testBufferPoolAlignment() throws {
|
||
let engine = try MarkBaseEngine()
|
||
|
||
let buf = engine.acquireBuffer(length: 100)
|
||
XCTAssertEqual(buf.length % 256, 0)
|
||
XCTAssertEqual(buf.length, 256)
|
||
}
|
||
|
||
func testBufferPoolReuse() throws {
|
||
let engine = try MarkBaseEngine()
|
||
|
||
for _ in 0..<10 {
|
||
let buf = engine.acquireBuffer(length: 2048)
|
||
engine.releaseBuffer(buf)
|
||
}
|
||
|
||
XCTAssertTrue(engine.bufferPool.totalReuses >= 9,
|
||
"Expected at least 9 reuses, got \(engine.bufferPool.totalReuses)")
|
||
}
|
||
|
||
func testBufferPoolStats() throws {
|
||
let engine = try MarkBaseEngine()
|
||
|
||
let buf1 = engine.acquireBuffer(length: 512)
|
||
let buf2 = engine.acquireBuffer(length: 1024)
|
||
engine.releaseBuffer(buf1)
|
||
let buf3 = engine.acquireBuffer(length: 512)
|
||
engine.releaseBuffer(buf2)
|
||
engine.releaseBuffer(buf3)
|
||
|
||
let stats = engine.bufferPool.stats
|
||
XCTAssertTrue(stats.contains("Allocations:"))
|
||
}
|
||
|
||
// ── Batch Dispatch Tests ──────────────────────────
|
||
|
||
func testBatchDispatch() throws {
|
||
let engine = try MarkBaseEngine()
|
||
try engine.compileSource(MetalKernels.vectorAdd)
|
||
let pipeline = try engine.pipeline(named: "vector_add")
|
||
|
||
let n: UInt = 128
|
||
let a = (0..<n).map { Float($0) }
|
||
let b = (0..<n).map { Float($0 * 2) }
|
||
|
||
let bufA = try engine.makeBuffer(a)
|
||
let bufB = try engine.makeBuffer(b)
|
||
let bufC = try engine.makeBuffer(length: Int(n) * MemoryLayout<Float>.stride)
|
||
let bufN = try engine.makeBuffer([n])
|
||
|
||
let gridSize = MTLSize(width: Int(n), height: 1, depth: 1)
|
||
|
||
try engine.batchDispatch([
|
||
(pipeline, [bufA, bufB, bufC, bufN], gridSize)
|
||
])
|
||
|
||
let result = engine.readFloats(from: bufC, count: Int(n))
|
||
for i in 0..<Int(n) {
|
||
XCTAssertEqual(result[i], a[i] + b[i], accuracy: 1e-5)
|
||
}
|
||
}
|
||
|
||
func testMakeBatchEncoder() throws {
|
||
let engine = try MarkBaseEngine()
|
||
try engine.compileSource(MetalKernels.vectorAdd)
|
||
let pipeline = try engine.pipeline(named: "vector_add")
|
||
|
||
let (cmdBuf, enc) = try engine.makeBatchEncoder()
|
||
|
||
enc.setComputePipelineState(pipeline)
|
||
let n: UInt = 64
|
||
let a = (0..<n).map { Float($0) }
|
||
let b = (0..<n).map { Float($0) }
|
||
|
||
let bufA = try engine.makeBuffer(a)
|
||
let bufB = try engine.makeBuffer(b)
|
||
let bufC = try engine.makeBuffer(length: Int(n) * MemoryLayout<Float>.stride)
|
||
let bufN = try engine.makeBuffer([n])
|
||
|
||
enc.setBuffer(bufA, offset: 0, index: 0)
|
||
enc.setBuffer(bufB, offset: 0, index: 1)
|
||
enc.setBuffer(bufC, offset: 0, index: 2)
|
||
enc.setBuffer(bufN, offset: 0, index: 3)
|
||
|
||
let tg = engine.threadgroupSize1D(pipeline, count: Int(n))
|
||
enc.dispatchThreads(MTLSize(width: Int(n), height: 1, depth: 1),
|
||
threadsPerThreadgroup: tg)
|
||
enc.endEncoding()
|
||
|
||
cmdBuf.commit()
|
||
cmdBuf.waitUntilCompleted()
|
||
|
||
let result = engine.readFloats(from: bufC, count: Int(n))
|
||
for i in 0..<Int(n) {
|
||
XCTAssertEqual(result[i], a[i] + b[i], accuracy: 1e-5)
|
||
}
|
||
}
|
||
|
||
// ── RMS Norm Chunked Tests ────────────────────────
|
||
|
||
func testRMSNormChunkedKernel() throws {
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
|
||
let size = 3840 // 12B hidden size
|
||
|
||
let input: [Float] = (0..<size).map { Float($0) / Float(size) }
|
||
let weight: [Float] = Array(repeating: 1.0, count: size)
|
||
|
||
let inputBuf = try engine.makeBuffer(input)
|
||
let weightBuf = try engine.makeBuffer(weight)
|
||
let outputBuf = try engine.makeBuffer(length: size * MemoryLayout<Float>.stride)
|
||
|
||
let pso = try engine.pipeline(named: "rms_norm_chunked")
|
||
let cmdBuf = engine.commandQueue.makeCommandBuffer()!
|
||
let enc = cmdBuf.makeComputeCommandEncoder()!
|
||
|
||
enc.setComputePipelineState(pso)
|
||
enc.setBuffer(inputBuf, offset: 0, index: 0)
|
||
enc.setBuffer(weightBuf, offset: 0, index: 1)
|
||
enc.setBuffer(outputBuf, offset: 0, index: 2)
|
||
var N = UInt32(size)
|
||
enc.setBytes(&N, length: MemoryLayout<UInt32>.size, index: 3)
|
||
var eps: Float = 1e-6
|
||
enc.setBytes(&eps, length: MemoryLayout<Float>.size, index: 4)
|
||
enc.setThreadgroupMemoryLength(256 * 4, index: 0)
|
||
|
||
let tg = MTLSize(width: 256, height: 1, depth: 1)
|
||
enc.dispatchThreads(MTLSize(width: size, height: 1, depth: 1),
|
||
threadsPerThreadgroup: tg)
|
||
enc.endEncoding()
|
||
cmdBuf.commit()
|
||
cmdBuf.waitUntilCompleted()
|
||
|
||
let result = engine.readFloats(from: outputBuf, count: size)
|
||
|
||
let sumSq = input.reduce(0) { $0 + $1 * $1 }
|
||
let rms = sqrt(sumSq / Float(size) + 1e-6)
|
||
|
||
for i in 0..<size {
|
||
let expected = input[i] / rms
|
||
XCTAssertEqual(result[i], expected, accuracy: 1e-3)
|
||
}
|
||
}
|
||
|
||
// ── HuggingFace Tokenizer Tests ───────────────────
|
||
|
||
func testHuggingFaceTokenizerRoundtrip() throws {
|
||
let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/main"
|
||
let tokenizerPath = modelDir + "/tokenizer.json"
|
||
|
||
guard FileManager.default.fileExists(atPath: tokenizerPath) else {
|
||
print("Skipping HuggingFaceTokenizer test - tokenizer.json not found")
|
||
return
|
||
}
|
||
|
||
let tokenizer = try HuggingFaceTokenizer(jsonPath: tokenizerPath)
|
||
|
||
// Verify vocab
|
||
XCTAssertGreaterThan(tokenizer.vocabSize, 0)
|
||
|
||
// Test roundtrip
|
||
let text = "Hello, world! This is a test."
|
||
let tokens = tokenizer.encode(text: text)
|
||
XCTAssertGreaterThan(tokens.count, 1) // At least BOS + tokens
|
||
|
||
let decoded = tokenizer.decode(tokens: tokens)
|
||
XCTAssertFalse(decoded.isEmpty)
|
||
|
||
print("Original: \(text)")
|
||
print("Tokens: \(tokens.count)")
|
||
print("Decoded: \(decoded)")
|
||
}
|
||
|
||
func testHuggingFaceTokenizerSpecialTokens() throws {
|
||
let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/main"
|
||
let tokenizerPath = modelDir + "/tokenizer.json"
|
||
|
||
guard FileManager.default.fileExists(atPath: tokenizerPath) else {
|
||
return
|
||
}
|
||
|
||
let tokenizer = try HuggingFaceTokenizer(jsonPath: tokenizerPath)
|
||
|
||
// Verify special token IDs
|
||
XCTAssertGreaterThanOrEqual(tokenizer.bosTokenId, 0)
|
||
XCTAssertGreaterThanOrEqual(tokenizer.eosTokenId, 0)
|
||
XCTAssertGreaterThanOrEqual(tokenizer.padTokenId, 0)
|
||
}
|
||
|
||
// ── Layer Config Tests ────────────────────────────
|
||
|
||
func test12BLayerConfig() {
|
||
// 12B should use nHeads=16, nKvHeads=8
|
||
let config = E4BLayerConfig.full(
|
||
hiddenSize: 3840,
|
||
headDim: 512,
|
||
intermediateSize: 30720,
|
||
nHeads: 16,
|
||
nKvHeads: 8
|
||
)
|
||
|
||
XCTAssertEqual(config.nHeads, 16)
|
||
XCTAssertEqual(config.nKvHeads, 8)
|
||
XCTAssertEqual(config.hiddenSize, 3840)
|
||
XCTAssertEqual(config.headDim, 512)
|
||
XCTAssertFalse(config.isSliding)
|
||
XCTAssertEqual(config.rotatedDim, 128) // 512 * 0.25
|
||
}
|
||
|
||
func test12BLayerConfigSliding() {
|
||
let config = E4BLayerConfig.sliding(
|
||
hiddenSize: 3840,
|
||
headDim: 256,
|
||
intermediateSize: 15360,
|
||
nHeads: 16,
|
||
nKvHeads: 8,
|
||
windowSize: 512
|
||
)
|
||
|
||
XCTAssertEqual(config.nHeads, 16)
|
||
XCTAssertEqual(config.nKvHeads, 8)
|
||
XCTAssertTrue(config.isSliding)
|
||
XCTAssertEqual(config.windowSize, 512)
|
||
}
|
||
|
||
// ── ForwardTemps 12B Tests ────────────────────────
|
||
|
||
func testForwardTemps12B() throws {
|
||
let engine = try MarkBaseEngine()
|
||
|
||
// Test 12B parameters
|
||
let temps = try ForwardTemps(
|
||
device: engine.device,
|
||
maxHeadDim: 512,
|
||
maxIntermediate: 30720, // 15360 * 2
|
||
hiddenSize: 3840,
|
||
nHeads: 16,
|
||
nKvHeads: 8
|
||
)
|
||
|
||
// Verify sizes
|
||
XCTAssertEqual(temps.q.length, 16 * 512 * 4, "q should be 16 heads × 512 dim")
|
||
XCTAssertEqual(temps.gate.length, 30720 * 4, "gate should be 30720")
|
||
XCTAssertEqual(temps.io.length, 3840 * 4, "io should be 3840")
|
||
}
|
||
}
|