v2: add debug prints to forwardLayer, identify hanging kernel

Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift

@@ -41,14 +41,12 @@ public struct EmbeddingGemmaConfig: Codable {
     }
 }
 
-/// EmbeddingGemma - Google's 300M parameter embedding model
 public final class EmbeddingGemmaModel: @unchecked Sendable {
     public let config: EmbeddingGemmaConfig
     public let engine: MarkBaseEngine
     public let tokenizer: Tokenizer
     public let reader: SafeTensorsReader
 
-    // GPU Buffers
     public var embedTokens: MTLBuffer!
     public var finalNorm: MTLBuffer!
     public var layerNorms: [[MTLBuffer]] = []
@@ -67,7 +65,6 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         self.config = try EmbeddingGemmaConfig.load(from: modelDir)
         self.tokenizer = try TokenizerFactory.load(modelDir: modelDir)
         self.reader = try SafeTensorsReader(path: modelDir + "/model.safetensors")
-
         try loadWeights()
         print("✓ EmbeddingGemma loaded (\(config.numHiddenLayers) layers, hidden=\(config.hiddenSize))")
     }
@@ -75,7 +72,6 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
     private func loadWeights() throws {
         let embedData = try readTensor("embed_tokens.weight")
         embedTokens = engine.device.makeBuffer(bytes: embedData, length: embedData.count * 4)!
-
         for i in 0..<config.numHiddenLayers {
             let p = "layers.\(i)"
             layerNorms.append([
@@ -94,49 +90,48 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
             upProjs.append(try loadBuffer("\(p).mlp.up_proj.weight"))
             downProjs.append(try loadBuffer("\(p).mlp.down_proj.weight"))
         }
-
         let fnData = try readTensor("norm.weight")
         finalNorm = engine.device.makeBuffer(bytes: fnData, length: fnData.count * 4)!
     }
 
-    /// Generate embedding for text
     public func embed(text: String, maxLen: Int = 2048) throws -> [Float] {
         var tokens = tokenizer.encode(text: text)
         if tokens.count > maxLen { tokens = Array(tokens.prefix(maxLen)) }
         guard !tokens.isEmpty else { return [] }
 
         let seqLen = tokens.count, hs = config.hiddenSize
-        print("  embed: seqLen=\(seqLen)")
 
-        // Use single command buffer for entire forward pass
+        // Test 1: Embedding lookup only
+        print("  TEST: Embedding lookup...")
+        let cmdBuf1 = engine.commandQueue.makeCommandBuffer()!
+        let inputBuf = try lookupEmbeddings(tokens: tokens, cmdBuf: cmdBuf1)
+        cmdBuf1.commit()
+        cmdBuf1.waitUntilCompleted()
+        print("  TEST: Embedding lookup OK")
+
+        // Test 2: Single layer forward (layer 0 only)
+        print("  TEST: Layer 0 forward...")
+        let cmdBuf2 = engine.commandQueue.makeCommandBuffer()!
+        var hidden = try forwardLayerDebug(hidden: inputBuf, layerIdx: 0, seqLen: seqLen, cmdBuf: cmdBuf2)
+        cmdBuf2.commit()
+        cmdBuf2.waitUntilCompleted()
+        print("  TEST: Layer 0 OK")
+
+        // Full forward pass
+        print("  TEST: Full forward pass...")
         let cmdBuf = engine.commandQueue.makeCommandBuffer()!
-
+        hidden = inputBuf
-        // Embedding lookup
-        print("  embed: lookupEmbeddings")
-        let inputBuf = try lookupEmbeddings(tokens: tokens, cmdBuf: cmdBuf)
-        print("  embed: lookupEmbeddings done")
-
-        // Forward through layers
-        var hidden = inputBuf
         for idx in 0..<config.numHiddenLayers {
-            if idx % 6 == 0 { print("  embed: layer \(idx)/\(config.numHiddenLayers)") }
+            if idx % 6 == 0 { print("  Layer \(idx)/\(config.numHiddenLayers)") }
-            hidden = try forwardLayer(hidden: hidden, layerIdx: idx, seqLen: seqLen, cmdBuf: cmdBuf)
+            hidden = try forwardLayerDebug(hidden: hidden, layerIdx: idx, seqLen: seqLen, cmdBuf: cmdBuf)
         }
-        print("  embed: all layers done")
+        print("  TEST: All layers OK")
 
-        // Final norm
-        print("  embed: final norm")
         let output = try applyRmsNorm(input: hidden, weight: finalNorm, count: seqLen * hs, cmdBuf: cmdBuf)
-        print("  embed: committing")
-
         cmdBuf.commit()
         cmdBuf.waitUntilCompleted()
-        print("  embed: completed")
 
-        // Readback
         let data = engine.readFloats(from: output, count: seqLen * hs)
-
-        // Mean pool + L2 normalize
         var embedding = [Float](repeating: 0, count: hs)
         for i in 0..<seqLen {
             let start = i * hs
@@ -144,29 +139,20 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         }
         let n = Float(seqLen)
         for i in 0..<hs { embedding[i] /= n }
-
         var norm: Float = 0
         for i in 0..<hs { norm += embedding[i] * embedding[i] }
         norm = sqrt(norm)
         if norm > 0 { for i in 0..<hs { embedding[i] /= norm } }
-
         return embedding
     }
 
-    // MARK: - Helpers (all use shared cmdBuf)
-
     private func readTensor(_ name: String) throws -> [Float] {
-        guard let desc = reader.tensor(named: name) else {
+        guard let desc = reader.tensor(named: name) else { throw WeightError.tensorNotFound(name) }
-            throw WeightError.tensorNotFound(name)
-        }
         let data = try reader.read(tensor: desc)
         switch desc.dtype {
-        case .f32:
+        case .f32: return data.withUnsafeBytes { Array(UnsafeBufferPointer(start: $0.baseAddress?.assumingMemoryBound(to: Float.self), count: data.count/4)) }
-            return data.withUnsafeBytes { Array(UnsafeBufferPointer(start: $0.baseAddress?.assumingMemoryBound(to: Float.self), count: data.count/4)) }
+        case .bf16: return try SafeTensorsReader.bf16ToFloat32(data)
-        case .bf16:
+        default: throw WeightError.unsupportedDtype(desc.dtype.rawValue)
-            return try SafeTensorsReader.bf16ToFloat32(data)
-        default:
-            throw WeightError.unsupportedDtype(desc.dtype.rawValue)
         }
     }
 
@@ -189,8 +175,7 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         enc.setBytes(&h, length: 4, index: 3)
         enc.setBytes(&s, length: 4, index: 4)
         enc.setBytes(&v, length: 4, index: 5)
-        enc.dispatchThreads(MTLSize(width: seqLen, height: 1, depth: 1),
+        enc.dispatchThreads(MTLSize(width: seqLen, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, seqLen), height: 1, depth: 1))
-                           threadsPerThreadgroup: MTLSize(width: min(256, seqLen), height: 1, depth: 1))
         return buf
     }
 
@@ -206,8 +191,7 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         var c = UInt32(count), e: Float = config.rmsNormEps
         enc.setBytes(&c, length: 4, index: 3)
         enc.setBytes(&e, length: 4, index: 4)
-        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
+        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
-                           threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
         return output
     }
 
@@ -226,8 +210,7 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
             enc.setBytes(&mm, length: 4, index: 3)
             enc.setBytes(&kk, length: 4, index: 4)
             enc.setBytes(&nn, length: 4, index: 5)
-            enc.dispatchThreads(MTLSize(width: n, height: 1, depth: 1),
+            enc.dispatchThreads(MTLSize(width: n, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, n), height: 1, depth: 1))
-                               threadsPerThreadgroup: MTLSize(width: min(256, n), height: 1, depth: 1))
         }
     }
 
@@ -241,8 +224,7 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         enc.setBuffer(output, offset: 0, index: 2)
         var c = UInt32(count)
         enc.setBytes(&c, length: 4, index: 3)
-        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
+        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
-                           threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
     }
 
     private func geluMul(gate: MTLBuffer, up: MTLBuffer, output: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws {
@@ -255,8 +237,7 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         enc.setBuffer(output, offset: 0, index: 2)
         var c = UInt32(count)
         enc.setBytes(&c, length: 4, index: 3)
-        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
+        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
-                           threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
     }
 
     private func applyRoPE(q: MTLBuffer, k: MTLBuffer, seqLen: Int, headDim: Int, numHeads: Int, cmdBuf: MTLCommandBuffer) throws {
@@ -272,8 +253,7 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         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),
+        enc.dispatchThreads(MTLSize(width: numHeads * headDim / 2, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, numHeads * headDim / 2), height: 1, depth: 1))
-                           threadsPerThreadgroup: MTLSize(width: min(256, numHeads * headDim / 2), height: 1, depth: 1))
     }
 
     private func bidirectionalAttention(q: MTLBuffer, k: MTLBuffer, v: MTLBuffer, output: MTLBuffer, seqLen: Int, cmdBuf: MTLCommandBuffer) throws {
@@ -296,71 +276,81 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         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),
+        enc.dispatchThreads(MTLSize(width: seqLen * config.numAttentionHeads, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, seqLen * config.numAttentionHeads), height: 1, depth: 1))
-                           threadsPerThreadgroup: MTLSize(width: min(256, seqLen * config.numAttentionHeads), height: 1, depth: 1))
     }
 
-    private func forwardLayer(hidden: MTLBuffer, layerIdx: Int, seqLen: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
+    private func forwardLayerDebug(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
+        print("    Residual copy...")
         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
+        print("    Input norm...")
         let h1 = try applyRmsNorm(input: hidden, weight: layerNorms[layerIdx][0], count: seqLen * hs, cmdBuf: cmdBuf)
 
-        // Q, K, V projections
+        print("    QKV 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)
+        print("    Q done")
         try matmulSeq(input: h1, weight: kProjs[layerIdx], output: kBuf, m: seqLen, k: hs, n: nKV * hDim, cmdBuf: cmdBuf)
+        print("    K done")
         try matmulSeq(input: h1, weight: vProjs[layerIdx], output: vBuf, m: seqLen, k: hs, n: nKV * hDim, cmdBuf: cmdBuf)
+        print("    V done")
 
-        // RoPE
+        print("    RoPE...")
         try applyRoPE(q: qBuf, k: kBuf, seqLen: seqLen, headDim: hDim, numHeads: nH, cmdBuf: cmdBuf)
+        print("    RoPE done")
 
-        // Bidirectional sliding window attention
+        print("    Attention...")
         let attnOut = device.makeBuffer(length: seqLen * nH * hDim * 4)!
         try bidirectionalAttention(q: qBuf, k: kBuf, v: vBuf, output: attnOut, seqLen: seqLen, cmdBuf: cmdBuf)
+        print("    Attention done")
 
-        // O projection
+        print("    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)
+        print("    O done")
 
-        // Post-attn norm
+        print("    Post-attn norm...")
         let h2n = try applyRmsNorm(input: h2, weight: layerNorms[layerIdx][2], count: seqLen * hs, cmdBuf: cmdBuf)
 
-        // Add residual: hidden = resid + h2n
+        print("    Add residual 1...")
         try eltwiseAdd(a: resid, b: h2n, output: hidden, count: seqLen * hs, cmdBuf: cmdBuf)
 
-        // Pre-FF norm
+        print("    Pre-FF norm...")
         let h3 = try applyRmsNorm(input: hidden, weight: layerNorms[layerIdx][1], count: seqLen * hs, cmdBuf: cmdBuf)
 
-        // MLP: gate, up
+        print("    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)
+        print("    Gate done")
         try matmulSeq(input: h3, weight: upProjs[layerIdx], output: up, m: seqLen, k: hs, n: intermedi, cmdBuf: cmdBuf)
+        print("    Up done")
 
-        // GELU(gate) * up
+        print("    GELU mul...")
         let gated = device.makeBuffer(length: seqLen * intermedi * 4)!
         try geluMul(gate: gate, up: up, output: gated, count: seqLen * intermedi, cmdBuf: cmdBuf)
+        print("    GELU done")
 
-        // Down projection
+        print("    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)
+        print("    Down done")
 
-        // Post-FF norm
+        print("    Post-FF norm...")
         let h4n = try applyRmsNorm(input: h4, weight: layerNorms[layerIdx][3], count: seqLen * hs, cmdBuf: cmdBuf)
 
-        // Add residual: hidden = hidden + h4n
+        print("    Add residual 2...")
         try eltwiseAdd(a: hidden, b: h4n, output: hidden, count: seqLen * hs, cmdBuf: cmdBuf)
+        print("    Layer \(layerIdx) complete")
 
         return hidden
     }