v2: add embedding tests, multilingual embedding support

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"
+    }
+}

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<Float>.size, index: 17)
         var topK = UInt32(topK)
         enc.setBytes(&topK, length: MemoryLayout<UInt32>.size, index: 18)
+        var groupSize = UInt32(expertGroupSize)
+        enc.setBytes(&groupSize, length: MemoryLayout<UInt32>.size, index: 19)
 
         let count = Int(max(hiddenSize, moeIntermediate))
         let logitStorage = Int(numExperts) + Int(topK) + Int(topK)

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;
 

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") {