diff --git a/Scripts/com.markbase.12b.plist b/Scripts/com.markbase.12b.plist new file mode 100644 index 0000000..620c795 --- /dev/null +++ b/Scripts/com.markbase.12b.plist @@ -0,0 +1,24 @@ + + + + + Label + com.markbase.12b + ProgramArguments + + /Users/accusys/MarkBaseEngine/.build/arm64-apple-macosx/release/MarkBaseServer + gemma-4-12b-it-4bit + 8081 + + RunAtLoad + + KeepAlive + + StandardOutPath + /Users/accusys/MarkBaseEngine/logs/12b.log + StandardErrorPath + /Users/accusys/MarkBaseEngine/logs/12b.log + WorkingDirectory + /Users/accusys/MarkBaseEngine + + diff --git a/Scripts/com.markbase.26b.plist b/Scripts/com.markbase.26b.plist new file mode 100644 index 0000000..c87cd9d --- /dev/null +++ b/Scripts/com.markbase.26b.plist @@ -0,0 +1,24 @@ + + + + + Label + com.markbase.26b + ProgramArguments + + /Users/accusys/MarkBaseEngine/.build/arm64-apple-macosx/release/MarkBaseServer + gemma-4-26b-standard + 8082 + + RunAtLoad + + KeepAlive + + StandardOutPath + /Users/accusys/MarkBaseEngine/logs/26b.log + StandardErrorPath + /Users/accusys/MarkBaseEngine/logs/26b.log + WorkingDirectory + /Users/accusys/MarkBaseEngine + + diff --git a/Scripts/com.markbase.31b.plist b/Scripts/com.markbase.31b.plist new file mode 100644 index 0000000..dd8a85c --- /dev/null +++ b/Scripts/com.markbase.31b.plist @@ -0,0 +1,24 @@ + + + + + Label + com.markbase.31b + ProgramArguments + + /Users/accusys/MarkBaseEngine/.build/arm64-apple-macosx/release/MarkBaseServer + gemma-4-31b-it-4bit + 8083 + + RunAtLoad + + KeepAlive + + StandardOutPath + /Users/accusys/MarkBaseEngine/logs/31b.log + StandardErrorPath + /Users/accusys/MarkBaseEngine/logs/31b.log + WorkingDirectory + /Users/accusys/MarkBaseEngine + + diff --git a/Scripts/com.markbase.e4b.plist b/Scripts/com.markbase.e4b.plist new file mode 100644 index 0000000..9ec5220 --- /dev/null +++ b/Scripts/com.markbase.e4b.plist @@ -0,0 +1,24 @@ + + + + + Label + com.markbase.e4b + ProgramArguments + + /Users/accusys/MarkBaseEngine/.build/arm64-apple-macosx/release/MarkBaseServer + E4B-MarkBase + 8080 + + RunAtLoad + + KeepAlive + + StandardOutPath + /Users/accusys/MarkBaseEngine/logs/e4b.log + StandardErrorPath + /Users/accusys/MarkBaseEngine/logs/e4b.log + WorkingDirectory + /Users/accusys/MarkBaseEngine + + diff --git a/Scripts/com.markbase.embedding.plist b/Scripts/com.markbase.embedding.plist new file mode 100644 index 0000000..703ac62 --- /dev/null +++ b/Scripts/com.markbase.embedding.plist @@ -0,0 +1,24 @@ + + + + + Label + com.markbase.embedding + ProgramArguments + + /Users/accusys/MarkBaseEngine/.build/arm64-apple-macosx/release/MarkBaseServer + E4B-MarkBase + 8084 + + RunAtLoad + + KeepAlive + + StandardOutPath + /Users/accusys/MarkBaseEngine/logs/embedding.log + StandardErrorPath + /Users/accusys/MarkBaseEngine/logs/embedding.log + WorkingDirectory + /Users/accusys/MarkBaseEngine + + diff --git a/Scripts/setup-servers.sh b/Scripts/setup-servers.sh new file mode 100755 index 0000000..cc2af9d --- /dev/null +++ b/Scripts/setup-servers.sh @@ -0,0 +1,53 @@ +#!/bin/bash +# Setup MarkBase model servers as launchd services + +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +LAUNCH_AGENTS="$HOME/Library/LaunchAgents" +LOG_DIR="/Users/accusys/MarkBaseEngine/logs" + +echo "Setting up MarkBase model servers..." + +# Create logs directory +mkdir -p "$LOG_DIR" + +# List of services to install +SERVICES=( + "com.markbase.e4b.plist" + "com.markbase.12b.plist" + "com.markbase.26b.plist" + "com.markbase.31b.plist" + "com.markbase.embedding.plist" +) + +for plist in "${SERVICES[@]}"; do + src="$SCRIPT_DIR/$plist" + dst="$LAUNCH_AGENTS/$plist" + + if [ ! -f "$src" ]; then + echo " SKIP: $plist not found" + continue + fi + + # Copy to LaunchAgents + cp "$src" "$dst" + echo " Installed: $plist" + + # Load the service (skip if already loaded) + label="${plist%.plist}" + if launchctl list | grep -q "$label"; then + echo " Reloading: $label" + launchctl unload "$dst" 2>/dev/null || true + fi + launchctl load "$dst" + echo " Loaded: $label" +done + +echo "" +echo "Done! Services:" +echo " E4B-MarkBase → http://localhost:8080" +echo " gemma-4-12b-it-4bit → http://localhost:8081" +echo " gemma-4-26b → http://localhost:8082" +echo " gemma-4-31b → http://localhost:8083" +echo " Embedding (E4B) → http://localhost:8084" diff --git a/Sources/MarkBase/Embedding/EmbeddingGemmaConfig.swift b/Sources/MarkBase/Embedding/EmbeddingGemmaConfig.swift new file mode 100644 index 0000000..013b9b6 --- /dev/null +++ b/Sources/MarkBase/Embedding/EmbeddingGemmaConfig.swift @@ -0,0 +1,34 @@ +import Foundation + +/// EmbeddingGemma model configuration +public struct EmbeddingGemmaConfig: Codable { + public let hiddenSize: Int + public let numHiddenLayers: Int + public let vocabSize: Int + public let numAttentionHeads: Int + public let numKeyValueHeads: Int + public let headDim: Int + public let intermediateSize: Int + public let maxPositionEmbeddings: Int + public let slidingWindow: Int + public let rmsNormEps: Float + public let ropeTheta: Float + public let useBidirectionalAttention: Bool + public let layerTypes: [String] + + enum CodingKeys: String, CodingKey { + case hiddenSize = "hidden_size" + case numHiddenLayers = "num_hidden_layers" + case vocabSize = "vocab_size" + case numAttentionHeads = "num_attention_heads" + case numKeyValueHeads = "num_key_value_heads" + case headDim = "head_dim" + case intermediateSize = "intermediate_size" + case maxPositionEmbeddings = "max_position_embeddings" + case slidingWindow = "sliding_window" + case rmsNormEps = "rms_norm_eps" + case ropeTheta = "rope_theta" + case useBidirectionalAttention = "use_bidirectional_attention" + case layerTypes = "layer_types" + } +} diff --git a/Sources/MarkBase/Layers/Layer.swift b/Sources/MarkBase/Layers/Layer.swift index 10e4387..b89f228 100644 --- a/Sources/MarkBase/Layers/Layer.swift +++ b/Sources/MarkBase/Layers/Layer.swift @@ -976,10 +976,8 @@ 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) @@ -1011,6 +1009,8 @@ func quantizedMatmulExpert(engine: MarkBaseEngine, cmdBuf: MTLCommandBuffer, enc.setBytes(&rScale, length: MemoryLayout.size, index: 17) var topK = UInt32(topK) enc.setBytes(&topK, length: MemoryLayout.size, index: 18) + var groupSize = UInt32(expertGroupSize) + enc.setBytes(&groupSize, length: MemoryLayout.size, index: 19) let count = Int(max(hiddenSize, moeIntermediate)) let logitStorage = Int(numExperts) + Int(topK) + Int(topK) diff --git a/Sources/MarkBase/Metal/OptimizedKernels.metal b/Sources/MarkBase/Metal/OptimizedKernels.metal index d00476a..6b1cdd7 100644 --- a/Sources/MarkBase/Metal/OptimizedKernels.metal +++ b/Sources/MarkBase/Metal/OptimizedKernels.metal @@ -815,13 +815,14 @@ kernel void moe_mega_kernel( constant uint &numExperts [[buffer(16)]], constant float &routerScale [[buffer(17)]], constant uint &topK [[buffer(18)]], + constant uint &groupSize [[buffer(19)]], threadgroup float *shared_space [[threadgroup(0)]], uint gid [[thread_position_in_grid]], uint tid [[thread_position_in_threadgroup]], uint tgSize [[threads_per_threadgroup]] ) { - uint numGroupsIn = hiddenSize / 64; - uint numGroupsOut = moeIntermediate / 64; + uint numGroupsIn = hiddenSize / groupSize; + uint numGroupsOut = moeIntermediate / groupSize; uint packedPerIn = hiddenSize / 8; uint packedPerOut = moeIntermediate / 8; @@ -841,10 +842,10 @@ kernel void moe_mega_kernel( for (uint g = 0; g < numGroupsIn; g++) { float scale = s_router[tid * numGroupsIn + g]; float bias = b_router[tid * numGroupsIn + g]; - uint wBase = tid * packedPerIn + g * 8; - uint xBase = g * 64; + uint wBase = tid * packedPerIn + g * (groupSize / 8); + uint xBase = g * groupSize; - for (uint p = 0; p < 8; p += 4) { + for (uint p = 0; p < groupSize / 8; p += 4) { device uint4 *rPtr = (device uint4*)(&w_router[wBase + p]); uint4 packed = *rPtr; @@ -971,10 +972,10 @@ kernel void moe_mega_kernel( float uScale = s_up[sUpBase + gid * numGroupsIn + g]; float uBias = b_up[sUpBase + gid * numGroupsIn + g]; - uint wb = gid * packedPerIn + g * 8; - uint xBase = g * 64; + uint wb = gid * packedPerIn + g * (groupSize / 8); + uint xBase = g * groupSize; - for (uint p = 0; p < 8; p += 4) { + for (uint p = 0; p < groupSize / 8; p += 4) { device uint4 *gPtr = (device uint4*)(&w_gate[wGateBase + wb + p]); device uint4 *uPtr = (device uint4*)(&w_up[wUpBase + wb + p]); uint4 gP = *gPtr; @@ -1047,10 +1048,10 @@ kernel void moe_mega_kernel( float scale = s_down[wDownBase + gid * numGroupsOut + g]; float bias = b_down[wDownBase + gid * numGroupsOut + g]; - uint wb = gid * packedPerOut + g * 8; - uint iBase = g * 64; + uint wb = gid * packedPerOut + g * (groupSize / 8); + uint iBase = g * groupSize; - for (uint p = 0; p < 8; p += 4) { + for (uint p = 0; p < groupSize / 8; p += 4) { device uint4 *wPtr = (device uint4*)(&w_down[wDownBase + wb + p]); uint4 packed = *wPtr; diff --git a/Sources/MarkBaseServer/SimpleServerApp.swift b/Sources/MarkBaseServer/SimpleServerApp.swift index 2f7836d..f0475ba 100644 --- a/Sources/MarkBaseServer/SimpleServerApp.swift +++ b/Sources/MarkBaseServer/SimpleServerApp.swift @@ -23,6 +23,7 @@ struct SimpleServerApp { let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 512) let tokenizer = try TokenizerFactory.load(modelDir: modelPath) let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine) + let embeddingModel = try TextEmbeddingModel(modelDir: modelPath, engine: engine, config: TextEmbeddingConfig()) print("✓ E4B loaded (\(model.numHiddenLayers) layers)") @@ -168,14 +169,15 @@ struct SimpleServerApp { "deployment": "docs/DEPLOYMENT.md", "performance": "docs/PERFORMANCE.md" }, - "notes": [ + "notes": [ "All responses are in JSON format", "Text generation only (multimodal not yet supported via API)", "E4B model with 42 layers, ~4B parameters", "For multimodal (vision/audio) support, use the MarkBase Swift library directly", "Streaming support is planned but not yet implemented", "Function calling uses native Gemma 4 special tokens", - "Messages can include tool_calls (assistant) and tool responses (tool role) for multi-turn function calling" + "Messages can include tool_calls (assistant) and tool responses (tool role) for multi-turn function calling", + "Text embeddings available via /v1/embeddings endpoint (OpenAI-compatible)" ] } """ @@ -287,6 +289,73 @@ struct SimpleServerApp { } } + router.post("/v1/embeddings") { request, _ in + let buffer = try await request.body.collect(upTo: .max) + let data = Data(buffer: buffer) + + guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any], + let input = json["input"] else { + return "{\"error\":\"invalid request\",\"type\":\"invalid_request_error\",\"code\":400,\"message\":\"missing 'input' field\"}" + } + + let modelId = (json["model"] as? String) ?? "e4b" + let encodingFormat = (json["encoding_format"] as? String) ?? "float" + + let inputs: [String] + if let str = input as? String { + inputs = [str] + } else if let arr = input as? [String] { + inputs = arr + } else { + return "{\"error\":\"invalid request\",\"type\":\"invalid_request_error\",\"code\":400,\"message\":\"'input' must be string or array of strings\"}" + } + + var embeddings: [[String: Any]] = [] + for (i, text) in inputs.enumerated() { + let t0 = Date() + let embedding = try embeddingModel.embed(text: text) + let duration = Date().timeIntervalSince(t0) + + let embeddingData: [String: Any] + if encodingFormat == "base64" { + let base64 = embedding.withUnsafeBytes { Data($0).base64EncodedString() } + embeddingData = [ + "object": "embedding", + "index": i, + "embedding": base64, + "usage_ms": Int(duration * 1000) + ] + } else { + embeddingData = [ + "object": "embedding", + "index": i, + "embedding": embedding, + "usage_ms": Int(duration * 1000) + ] + } + embeddings.append(embeddingData) + } + + let id = UUID().uuidString + let ts = Int(Date().timeIntervalSince1970) + let totalTokens = inputs.reduce(0) { $0 + tokenizer.encode(text: $1).count } + + let response: [String: Any] = [ + "id": id, + "object": "list", + "created": ts, + "model": modelId, + "data": embeddings, + "usage": [ + "prompt_tokens": totalTokens, + "total_tokens": totalTokens + ] + ] + + let jsonData = try JSONSerialization.data(withJSONObject: response) + return String(data: jsonData, encoding: .utf8) ?? "{}" + } + let app = Application( router: router, configuration: .init(address: .hostname("0.0.0.0", port: port)) @@ -299,6 +368,7 @@ struct SimpleServerApp { print(" GET /health - Health check") print(" GET /v1/models - Model list") print(" POST /v1/chat/completions - Chat completion") + print(" POST /v1/embeddings - Text embeddings") print("") print("Model: \(modelName)") if modelName.contains("E4B") { diff --git a/Tests/01_Model/EmbeddingTest.swift b/Tests/01_Model/EmbeddingTest.swift new file mode 100644 index 0000000..684d0c1 --- /dev/null +++ b/Tests/01_Model/EmbeddingTest.swift @@ -0,0 +1,80 @@ +import XCTest +@testable import MarkBase + +final class EmbeddingTest: XCTestCase { + var engine: MarkBaseEngine! + var model: E4BModel! + var embeddingModel: TextEmbeddingModel! + let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase" + + override func setUp() { + super.setUp() + guard FileManager.default.fileExists(atPath: modelDir + "/model.safetensors") else { + return + } + engine = try? MarkBaseEngine(autoCompile: true) + model = try? E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512) + embeddingModel = try? TextEmbeddingModel(modelDir: modelDir, engine: engine, config: TextEmbeddingConfig()) + } + + func testEmbeddingDimension() throws { + try XCTSkipIf(embeddingModel == nil, "Model not found") + let embedding = try embeddingModel.embed(text: "Hello world") + XCTAssertEqual(embedding.count, 2560, "Embedding dimension should be 2560") + } + + func testEmbeddingNormalized() throws { + try XCTSkipIf(embeddingModel == nil, "Model not found") + let embedding = try embeddingModel.embed(text: "Test text") + let norm = sqrt(embedding.reduce(0) { $0 + $1 * $1 }) + XCTAssertEqual(norm, 1.0, accuracy: 0.001, "Embedding should be L2 normalized") + } + + func testSimilarSentences() throws { + try XCTSkipIf(embeddingModel == nil, "Model not found") + let e1 = try embeddingModel.embed(text: "The cat is sitting on the mat") + let e2 = try embeddingModel.embed(text: "A cat rests on a rug") + let e3 = try embeddingModel.embed(text: "The stock market crashed today") + + let sim12 = cosineSimilarity(e1, e2) + let sim13 = cosineSimilarity(e1, e3) + print("Similar(cat, cat): \(sim12)") + print("Similar(cat, stock): \(sim13)") + XCTAssertGreaterThan(sim12, sim13, "Similar sentences should have higher cosine similarity") + } + + func testDifferentLengths() throws { + try XCTSkipIf(embeddingModel == nil, "Model not found") + let e1 = try embeddingModel.embed(text: "Hi") + let e2 = try embeddingModel.embed(text: "This is a much longer sentence with many words") + XCTAssertEqual(e1.count, e2.count, "Embeddings should have same dimension regardless of input length") + XCTAssertEqual(e1.count, 2560) + } + + func testEmptyInput() throws { + try XCTSkipIf(embeddingModel == nil, "Model not found") + let embedding = try embeddingModel.embed(text: "") + XCTAssertEqual(embedding.count, 0, "Empty input should return empty embedding") + } + + func testBatchEmbedding() throws { + try XCTSkipIf(embeddingModel == nil, "Model not found") + let texts = ["Hello", "World", "Test"] + let embeddings = try embeddingModel.embedBatch(texts: texts) + XCTAssertEqual(embeddings.count, 3) + for embedding in embeddings { + XCTAssertEqual(embedding.count, 2560) + } + } + + private func cosineSimilarity(_ a: [Float], _ b: [Float]) -> Float { + guard a.count == b.count, !a.isEmpty else { return 0 } + var dot: Float = 0, normA: Float = 0, normB: Float = 0 + for i in 0.. Same-lingual different topic: \(simSame > simDiff)") + } + + private func cosineSimilarity(_ a: [Float], _ b: [Float]) -> Float { + guard a.count == b.count, !a.isEmpty else { return 0 } + var dot: Float = 0, normA: Float = 0, normB: Float = 0 + for i in 0..