diff --git a/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift b/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift index 9e55d4f..2467bfb 100644 --- a/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift +++ b/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift @@ -2,7 +2,7 @@ import Foundation import Metal import Accelerate -/// EmbeddingGemmaConfig - Configuration for EmbeddingGemma model +/// EmbeddingGemma configuration public struct EmbeddingGemmaConfig: Codable { public let hiddenSize: Int public let numHiddenLayers: Int @@ -73,12 +73,6 @@ public final class EmbeddingGemmaModel: @unchecked Sendable { } private func loadWeights() throws { - let hs = config.hiddenSize - let intermedi = config.intermediateSize - let nKV = config.numKeyValueHeads - let hDim = config.headDim - - // Embedding table [vocab, hidden] let embedData = try readTensor("embed_tokens.weight") embedTokens = engine.device.makeBuffer(bytes: embedData, length: embedData.count * 4)! @@ -90,15 +84,15 @@ public final class EmbeddingGemmaModel: @unchecked Sendable { try loadBuffer("\(p).post_attention_layernorm.weight"), try loadBuffer("\(p).post_feedforward_layernorm.weight"), ]) - qProjs.append(try loadBuffer("\(p).self_attn.q_proj.weight")) // [hs, hs] - kProjs.append(try loadBuffer("\(p).self_attn.k_proj.weight")) // [nKV*hDim, hs] - vProjs.append(try loadBuffer("\(p).self_attn.v_proj.weight")) // [nKV*hDim, hs] - oProjs.append(try loadBuffer("\(p).self_attn.o_proj.weight")) // [hs, nH*hDim] - qNorms.append(try loadBuffer("\(p).self_attn.q_norm.weight")) // [hDim] - kNorms.append(try loadBuffer("\(p).self_attn.k_norm.weight")) // [hDim] - gateProjs.append(try loadBuffer("\(p).mlp.gate_proj.weight")) // [intermedi, hs] - upProjs.append(try loadBuffer("\(p).mlp.up_proj.weight")) // [intermedi, hs] - downProjs.append(try loadBuffer("\(p).mlp.down_proj.weight")) // [hs, intermedi] + qProjs.append(try loadBuffer("\(p).self_attn.q_proj.weight")) + kProjs.append(try loadBuffer("\(p).self_attn.k_proj.weight")) + vProjs.append(try loadBuffer("\(p).self_attn.v_proj.weight")) + oProjs.append(try loadBuffer("\(p).self_attn.o_proj.weight")) + qNorms.append(try loadBuffer("\(p).self_attn.q_norm.weight")) + kNorms.append(try loadBuffer("\(p).self_attn.k_norm.weight")) + gateProjs.append(try loadBuffer("\(p).mlp.gate_proj.weight")) + upProjs.append(try loadBuffer("\(p).mlp.up_proj.weight")) + downProjs.append(try loadBuffer("\(p).mlp.down_proj.weight")) } let fnData = try readTensor("norm.weight") @@ -113,17 +107,23 @@ public final class EmbeddingGemmaModel: @unchecked Sendable { let seqLen = tokens.count, hs = config.hiddenSize + // Use single command buffer for entire forward pass + let cmdBuf = engine.commandQueue.makeCommandBuffer()! + // Embedding lookup - let inputBuf = try lookupEmbeddings(tokens: tokens) + let inputBuf = try lookupEmbeddings(tokens: tokens, cmdBuf: cmdBuf) // Forward through layers var hidden = inputBuf for idx in 0..