v2: add embedding tests, multilingual embedding support

2026-07-06 08:01:52 +08:00
parent ba4c41c29f
commit 85dd87e28a
13 changed files with 484 additions and 15 deletions
@@ -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)