v2: EmbeddingGemma - single cmdBuf fix, needs Metal kernel compilation

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]
+            qProjs.append(try loadBuffer("\(p).self_attn.q_proj.weight"))
-            kProjs.append(try loadBuffer("\(p).self_attn.k_proj.weight"))     // [nKV*hDim, hs]
+            kProjs.append(try loadBuffer("\(p).self_attn.k_proj.weight"))
-            vProjs.append(try loadBuffer("\(p).self_attn.v_proj.weight"))     // [nKV*hDim, hs]
+            vProjs.append(try loadBuffer("\(p).self_attn.v_proj.weight"))
-            oProjs.append(try loadBuffer("\(p).self_attn.o_proj.weight"))     // [hs, nH*hDim]
+            oProjs.append(try loadBuffer("\(p).self_attn.o_proj.weight"))
-            qNorms.append(try loadBuffer("\(p).self_attn.q_norm.weight"))     // [hDim]
+            qNorms.append(try loadBuffer("\(p).self_attn.q_norm.weight"))
-            kNorms.append(try loadBuffer("\(p).self_attn.k_norm.weight"))     // [hDim]
+            kNorms.append(try loadBuffer("\(p).self_attn.k_norm.weight"))
-            gateProjs.append(try loadBuffer("\(p).mlp.gate_proj.weight"))     // [intermedi, hs]
+            gateProjs.append(try loadBuffer("\(p).mlp.gate_proj.weight"))
-            upProjs.append(try loadBuffer("\(p).mlp.up_proj.weight"))         // [intermedi, hs]
+            upProjs.append(try loadBuffer("\(p).mlp.up_proj.weight"))
-            downProjs.append(try loadBuffer("\(p).mlp.down_proj.weight"))     // [hs, intermedi]
+            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..<config.numHiddenLayers {
-            hidden = try forwardLayer(hidden: hidden, layerIdx: idx, seqLen: seqLen)
+            hidden = try forwardLayer(hidden: hidden, layerIdx: idx, seqLen: seqLen, cmdBuf: cmdBuf)
         }
 
         // Final norm
-        let output = try applyRmsNorm(input: hidden, weight: finalNorm, count: seqLen * hs)
+        let output = try applyRmsNorm(input: hidden, weight: finalNorm, count: seqLen * hs, cmdBuf: cmdBuf)
+
+        cmdBuf.commit()
+        cmdBuf.waitUntilCompleted()
 
         // Readback
         let data = engine.readFloats(from: output, count: seqLen * hs)
@@ -145,7 +145,7 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         return embedding
     }
 
-    // MARK: - Helpers
+    // MARK: - Helpers (all use shared cmdBuf)
 
     private func readTensor(_ name: String) throws -> [Float] {
         guard let desc = reader.tensor(named: name) else {
@@ -167,10 +167,9 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         return engine.device.makeBuffer(bytes: data, length: data.count * 4)!
     }
 
-    private func lookupEmbeddings(tokens: [Int]) throws -> MTLBuffer {
+    private func lookupEmbeddings(tokens: [Int], cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
         let seqLen = tokens.count, hs = config.hiddenSize
         let buf = engine.device.makeBuffer(length: seqLen * hs * 4)!
-        let cmdBuf = engine.commandQueue.makeCommandBuffer()!
         let enc = cmdBuf.makeComputeCommandEncoder()!
         let pso = try engine.pipeline(named: "lookup_embeddings")
         enc.setComputePipelineState(pso)
@@ -184,13 +183,11 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         enc.dispatchThreads(MTLSize(width: seqLen, height: 1, depth: 1),
                            threadsPerThreadgroup: MTLSize(width: min(256, seqLen), height: 1, depth: 1))
         enc.endEncoding()
-        cmdBuf.commit(); cmdBuf.waitUntilCompleted()
         return buf
     }
 
-    private func applyRmsNorm(input: MTLBuffer, weight: MTLBuffer, count: Int) throws -> MTLBuffer {
+    private func applyRmsNorm(input: MTLBuffer, weight: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
         let output = engine.device.makeBuffer(length: count * 4)!
-        let cmdBuf = engine.commandQueue.makeCommandBuffer()!
         let enc = cmdBuf.makeComputeCommandEncoder()!
         let pso = try engine.pipeline(named: "rms_norm")
         enc.setComputePipelineState(pso)
@@ -203,89 +200,13 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
                            threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
         enc.endEncoding()
-        cmdBuf.commit(); cmdBuf.waitUntilCompleted()
         return output
     }
 
-    private func forwardLayer(hidden: MTLBuffer, layerIdx: Int, seqLen: Int) throws -> MTLBuffer {
-        let hs = config.hiddenSize, device = engine.device
-        let hDim = config.headDim, nH = config.numAttentionHeads, nKV = config.numKeyValueHeads
-        let intermedi = config.intermediateSize
-
-        let cmdBuf = engine.commandQueue.makeCommandBuffer()!
-
-        // Residual
-        let resid = device.makeBuffer(length: seqLen * hs * 4)!
-        let blit = cmdBuf.makeBlitCommandEncoder()!
-        blit.copy(from: hidden, sourceOffset: 0, to: resid, destinationOffset: 0, size: seqLen * hs * 4)
-        blit.endEncoding()
-
-        // Input norm
-        let h1 = try applyRmsNorm(input: hidden, weight: layerNorms[layerIdx][0], count: seqLen * hs)
-
-        // Q, K, V projections (using optimized matmul)
-        let qBuf = device.makeBuffer(length: seqLen * nH * hDim * 4)!
-        let kBuf = device.makeBuffer(length: seqLen * nKV * hDim * 4)!
-        let vBuf = device.makeBuffer(length: seqLen * nKV * hDim * 4)!
-        try matmulSeq(input: h1, weight: qProjs[layerIdx], output: qBuf, m: seqLen, k: hs, n: nH * hDim, cmdBuf: cmdBuf)
-        try matmulSeq(input: h1, weight: kProjs[layerIdx], output: kBuf, m: seqLen, k: hs, n: nKV * hDim, cmdBuf: cmdBuf)
-        try matmulSeq(input: h1, weight: vProjs[layerIdx], output: vBuf, m: seqLen, k: hs, n: nKV * hDim, cmdBuf: cmdBuf)
-
-        // RoPE
-        try applyRoPE(q: qBuf, k: kBuf, seqLen: seqLen, headDim: hDim, numHeads: nH, numKVHeads: nKV, cmdBuf: cmdBuf)
-
-        // Q/K Norm
-        try applyQKNorm(q: qBuf, k: kBuf, qNorm: qNorms[layerIdx], kNorm: kNorms[layerIdx], seqLen: seqLen, headDim: hDim, numHeads: nH, numKVHeads: nKV, cmdBuf: cmdBuf)
-
-        // Bidirectional sliding window attention
-        let attnOut = device.makeBuffer(length: seqLen * nH * hDim * 4)!
-        try bidirectionalAttention(q: qBuf, k: kBuf, v: vBuf, output: attnOut, seqLen: seqLen, cmdBuf: cmdBuf)
-
-        // O projection
-        let h2 = device.makeBuffer(length: seqLen * hs * 4)!
-        try matmulSeq(input: attnOut, weight: oProjs[layerIdx], output: h2, m: seqLen, k: nH * hDim, n: hs, cmdBuf: cmdBuf)
-
-        // Post-attn norm
-        let h2n = try applyRmsNorm(input: h2, weight: layerNorms[layerIdx][2], count: seqLen * hs)
-
-        // Add residual: hidden = resid + h2n
-        try eltwiseAdd(a: resid, b: h2n, output: hidden, count: seqLen * hs, cmdBuf: cmdBuf)
-
-        // Pre-FF norm
-        let h3 = try applyRmsNorm(input: hidden, weight: layerNorms[layerIdx][1], count: seqLen * hs)
-
-        // MLP: gate, up
-        let gate = device.makeBuffer(length: seqLen * intermedi * 4)!
-        let up = device.makeBuffer(length: seqLen * intermedi * 4)!
-        try matmulSeq(input: h3, weight: gateProjs[layerIdx], output: gate, m: seqLen, k: hs, n: intermedi, cmdBuf: cmdBuf)
-        try matmulSeq(input: h3, weight: upProjs[layerIdx], output: up, m: seqLen, k: hs, n: intermedi, cmdBuf: cmdBuf)
-
-        // GELU(gate) * up
-        let gated = device.makeBuffer(length: seqLen * intermedi * 4)!
-        try geluMul(gate: gate, up: up, output: gated, count: seqLen * intermedi, cmdBuf: cmdBuf)
-
-        // Down projection
-        let h4 = device.makeBuffer(length: seqLen * hs * 4)!
-        try matmulSeq(input: gated, weight: downProjs[layerIdx], output: h4, m: seqLen, k: intermedi, n: hs, cmdBuf: cmdBuf)
-
-        // Post-FF norm
-        let h4n = try applyRmsNorm(input: h4, weight: layerNorms[layerIdx][3], count: seqLen * hs)
-
-        // Add residual: hidden = hidden + h4n
-        try eltwiseAdd(a: hidden, b: h4n, output: hidden, count: seqLen * hs, cmdBuf: cmdBuf)
-
-        cmdBuf.commit(); cmdBuf.waitUntilCompleted()
-        return hidden
-    }
-
-    // MARK: - Metal Kernels
-
     private func matmulSeq(input: MTLBuffer, weight: MTLBuffer, output: MTLBuffer, m: Int, k: Int, n: Int, cmdBuf: MTLCommandBuffer) throws {
         let enc = cmdBuf.makeComputeCommandEncoder()!
         let pso = try engine.pipeline(named: "matmul_f32")
         enc.setComputePipelineState(pso)
-        // Note: matmul_f32 only handles M=1, so we loop over positions
-        // For production: implement a proper batch matmul kernel
         for i in 0..<m {
             let inputOffset = i * k * 4
             let outputOffset = i * n * 4
@@ -302,52 +223,6 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         enc.endEncoding()
     }
 
-    private func applyRoPE(q: MTLBuffer, k: MTLBuffer, seqLen: Int, headDim: Int, numHeads: Int, numKVHeads: Int, cmdBuf: MTLCommandBuffer) throws {
-        let enc = cmdBuf.makeComputeCommandEncoder()!
-        let pso = try engine.pipeline(named: "apply_rope")
-        enc.setComputePipelineState(pso)
-        enc.setBuffer(q, offset: 0, index: 0)
-        enc.setBuffer(k, offset: 0, index: 1)
-        var sl = UInt32(seqLen), hd = UInt32(headDim), nh = UInt32(numHeads)
-        var rt: Float = Float(config.ropeTheta)
-        enc.setBytes(&sl, length: 4, index: 2)
-        enc.setBytes(&hd, length: 4, index: 3)
-        enc.setBytes(&nh, length: 4, index: 4)
-        enc.setBytes(&rt, length: 4, index: 5)
-        enc.dispatchThreads(MTLSize(width: numHeads * headDim / 2, height: 1, depth: 1),
-                           threadsPerThreadgroup: MTLSize(width: min(256, numHeads * headDim / 2), height: 1, depth: 1))
-        enc.endEncoding()
-    }
-
-    private func applyQKNorm(q: MTLBuffer, k: MTLBuffer, qNorm: MTLBuffer, kNorm: MTLBuffer, seqLen: Int, headDim: Int, numHeads: Int, numKVHeads: Int, cmdBuf: MTLCommandBuffer) throws {
-        // Apply RMSNorm per head
-        // TODO: Implement per-head normalization
-    }
-
-    private func bidirectionalAttention(q: MTLBuffer, k: MTLBuffer, v: MTLBuffer, output: MTLBuffer, seqLen: Int, cmdBuf: MTLCommandBuffer) throws {
-        let enc = cmdBuf.makeComputeCommandEncoder()!
-        let pso = try engine.pipeline(named: "bidirectional_sliding_attn")
-        enc.setComputePipelineState(pso)
-        enc.setBuffer(q, offset: 0, index: 0)
-        enc.setBuffer(k, offset: 0, index: 1)
-        enc.setBuffer(v, offset: 0, index: 2)
-        enc.setBuffer(output, offset: 0, index: 3)
-        var sl = UInt32(seqLen), hd = UInt32(config.headDim), nh = UInt32(config.numAttentionHeads)
-        var nkv = UInt32(config.numKeyValueHeads), sw = UInt32(config.slidingWindow)
-        var scale: Float = 1.0 / sqrt(Float(config.headDim))
-        enc.setBytes(&sl, length: 4, index: 4)
-        enc.setBytes(&hd, length: 4, index: 5)
-        enc.setBytes(&nh, length: 4, index: 6)
-        enc.setBytes(&nkv, length: 4, index: 7)
-        enc.setBytes(&sw, length: 4, index: 8)
-        enc.setBytes(&scale, length: 4, index: 9)
-        let tgMem = config.slidingWindow * 4  // shared memory for scores
-        enc.setThreadgroupMemoryLength(tgMem, index: 0)
-        enc.dispatchThreads(MTLSize(width: seqLen * config.numAttentionHeads, height: 1, depth: 1),
-                           threadsPerThreadgroup: MTLSize(width: min(256, seqLen * config.numAttentionHeads), height: 1, depth: 1))
-        enc.endEncoding()
-    }
-
     private func eltwiseAdd(a: MTLBuffer, b: MTLBuffer, output: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws {
         let enc = cmdBuf.makeComputeCommandEncoder()!
         let pso = try engine.pipeline(named: "eltwise_add")
@@ -375,4 +250,110 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
                            threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
         enc.endEncoding()
     }
+
+    private func applyRoPE(q: MTLBuffer, k: MTLBuffer, seqLen: Int, headDim: Int, numHeads: Int, cmdBuf: MTLCommandBuffer) throws {
+        let enc = cmdBuf.makeComputeCommandEncoder()!
+        let pso = try engine.pipeline(named: "apply_rope")
+        enc.setComputePipelineState(pso)
+        enc.setBuffer(q, offset: 0, index: 0)
+        enc.setBuffer(k, offset: 0, index: 1)
+        var sl = UInt32(seqLen), hd = UInt32(headDim), nh = UInt32(numHeads)
+        var rt: Float = Float(config.ropeTheta)
+        enc.setBytes(&sl, length: 4, index: 2)
+        enc.setBytes(&hd, length: 4, index: 3)
+        enc.setBytes(&nh, length: 4, index: 4)
+        enc.setBytes(&rt, length: 4, index: 5)
+        enc.dispatchThreads(MTLSize(width: numHeads * headDim / 2, height: 1, depth: 1),
+                           threadsPerThreadgroup: MTLSize(width: min(256, numHeads * headDim / 2), height: 1, depth: 1))
+        enc.endEncoding()
+    }
+
+    private func bidirectionalAttention(q: MTLBuffer, k: MTLBuffer, v: MTLBuffer, output: MTLBuffer, seqLen: Int, cmdBuf: MTLCommandBuffer) throws {
+        let enc = cmdBuf.makeComputeCommandEncoder()!
+        let pso = try engine.pipeline(named: "bidirectional_sliding_attn")
+        enc.setComputePipelineState(pso)
+        enc.setBuffer(q, offset: 0, index: 0)
+        enc.setBuffer(k, offset: 0, index: 1)
+        enc.setBuffer(v, offset: 0, index: 2)
+        enc.setBuffer(output, offset: 0, index: 3)
+        var sl = UInt32(seqLen), hd = UInt32(config.headDim), nh = UInt32(config.numAttentionHeads)
+        var nkv = UInt32(config.numKeyValueHeads), sw = UInt32(config.slidingWindow)
+        var scale: Float = 1.0 / sqrt(Float(config.headDim))
+        enc.setBytes(&sl, length: 4, index: 4)
+        enc.setBytes(&hd, length: 4, index: 5)
+        enc.setBytes(&nh, length: 4, index: 6)
+        enc.setBytes(&nkv, length: 4, index: 7)
+        enc.setBytes(&sw, length: 4, index: 8)
+        enc.setBytes(&scale, length: 4, index: 9)
+        let tgMem = config.slidingWindow * 4
+        enc.setThreadgroupMemoryLength(tgMem, index: 0)
+        enc.dispatchThreads(MTLSize(width: seqLen * config.numAttentionHeads, height: 1, depth: 1),
+                           threadsPerThreadgroup: MTLSize(width: min(256, seqLen * config.numAttentionHeads), height: 1, depth: 1))
+        enc.endEncoding()
+    }
+
+    private func forwardLayer(hidden: MTLBuffer, layerIdx: Int, seqLen: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
+        let hs = config.hiddenSize, device = engine.device
+        let hDim = config.headDim, nH = config.numAttentionHeads, nKV = config.numKeyValueHeads
+        let intermedi = config.intermediateSize
+
+        // Residual
+        let resid = device.makeBuffer(length: seqLen * hs * 4)!
+        let blit = cmdBuf.makeBlitCommandEncoder()!
+        blit.copy(from: hidden, sourceOffset: 0, to: resid, destinationOffset: 0, size: seqLen * hs * 4)
+        blit.endEncoding()
+
+        // Input norm
+        let h1 = try applyRmsNorm(input: hidden, weight: layerNorms[layerIdx][0], count: seqLen * hs, cmdBuf: cmdBuf)
+
+        // Q, K, V projections
+        let qBuf = device.makeBuffer(length: seqLen * nH * hDim * 4)!
+        let kBuf = device.makeBuffer(length: seqLen * nKV * hDim * 4)!
+        let vBuf = device.makeBuffer(length: seqLen * nKV * hDim * 4)!
+        try matmulSeq(input: h1, weight: qProjs[layerIdx], output: qBuf, m: seqLen, k: hs, n: nH * hDim, cmdBuf: cmdBuf)
+        try matmulSeq(input: h1, weight: kProjs[layerIdx], output: kBuf, m: seqLen, k: hs, n: nKV * hDim, cmdBuf: cmdBuf)
+        try matmulSeq(input: h1, weight: vProjs[layerIdx], output: vBuf, m: seqLen, k: hs, n: nKV * hDim, cmdBuf: cmdBuf)
+
+        // RoPE
+        try applyRoPE(q: qBuf, k: kBuf, seqLen: seqLen, headDim: hDim, numHeads: nH, cmdBuf: cmdBuf)
+
+        // Bidirectional sliding window attention
+        let attnOut = device.makeBuffer(length: seqLen * nH * hDim * 4)!
+        try bidirectionalAttention(q: qBuf, k: kBuf, v: vBuf, output: attnOut, seqLen: seqLen, cmdBuf: cmdBuf)
+
+        // O projection
+        let h2 = device.makeBuffer(length: seqLen * hs * 4)!
+        try matmulSeq(input: attnOut, weight: oProjs[layerIdx], output: h2, m: seqLen, k: nH * hDim, n: hs, cmdBuf: cmdBuf)
+
+        // Post-attn norm
+        let h2n = try applyRmsNorm(input: h2, weight: layerNorms[layerIdx][2], count: seqLen * hs, cmdBuf: cmdBuf)
+
+        // Add residual: hidden = resid + h2n
+        try eltwiseAdd(a: resid, b: h2n, output: hidden, count: seqLen * hs, cmdBuf: cmdBuf)
+
+        // Pre-FF norm
+        let h3 = try applyRmsNorm(input: hidden, weight: layerNorms[layerIdx][1], count: seqLen * hs, cmdBuf: cmdBuf)
+
+        // MLP: gate, up
+        let gate = device.makeBuffer(length: seqLen * intermedi * 4)!
+        let up = device.makeBuffer(length: seqLen * intermedi * 4)!
+        try matmulSeq(input: h3, weight: gateProjs[layerIdx], output: gate, m: seqLen, k: hs, n: intermedi, cmdBuf: cmdBuf)
+        try matmulSeq(input: h3, weight: upProjs[layerIdx], output: up, m: seqLen, k: hs, n: intermedi, cmdBuf: cmdBuf)
+
+        // GELU(gate) * up
+        let gated = device.makeBuffer(length: seqLen * intermedi * 4)!
+        try geluMul(gate: gate, up: up, output: gated, count: seqLen * intermedi, cmdBuf: cmdBuf)
+
+        // Down projection
+        let h4 = device.makeBuffer(length: seqLen * hs * 4)!
+        try matmulSeq(input: gated, weight: downProjs[layerIdx], output: h4, m: seqLen, k: intermedi, n: hs, cmdBuf: cmdBuf)
+
+        // Post-FF norm
+        let h4n = try applyRmsNorm(input: h4, weight: layerNorms[layerIdx][3], count: seqLen * hs, cmdBuf: cmdBuf)
+
+        // Add residual: hidden = hidden + h4n
+        try eltwiseAdd(a: hidden, b: h4n, output: hidden, count: seqLen * hs, cmdBuf: cmdBuf)
+
+        return hidden
+    }
 }

Tests/01_Model/EmbeddingGemmaTest.swift

@@ -0,0 +1,29 @@
+import XCTest
+@testable import MarkBase
+
+final class EmbeddingGemmaTest: XCTestCase {
+    func testLoadAndEmbed() throws {
+        let modelDir = "/Users/accusys/MarkBaseEngine/models/embeddinggemma-300m"
+        guard FileManager.default.fileExists(atPath: modelDir + "/model.safetensors") else {
+            XCTFail("Model not found"); return
+        }
+        let engine = try MarkBaseEngine(autoCompile: true)
+        let model = try EmbeddingGemmaModel(modelDir: modelDir, engine: engine)
+        XCTAssertEqual(model.config.numHiddenLayers, 24)
+        XCTAssertEqual(model.config.hiddenSize, 768)
+        
+        let emb = try model.embed(text: "Hello world")
+        XCTAssertEqual(emb.count, 768)
+        
+        // Check L2 norm is 1.0
+        var norm: Float = 0
+        for v in emb { norm += v * v }
+        norm = sqrt(norm)
+        XCTAssertEqual(norm, 1.0, accuracy: 0.001, "L2 norm should be 1.0")
+        
+        // Check no NaN
+        XCTAssertFalse(emb.contains { $0.isNaN }, "No NaN values")
+        
+        print("✓ Embedding dim=\(emb.count), norm=\(String(format: "%.4f", norm))")
+    }
+}