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;