From 492b779634c74710a0c854b889307488e8671bca Mon Sep 17 00:00:00 2001
From: MarkBase Admin <admin@markbase.local>
Date: Mon, 6 Jul 2026 13:05:30 +0800
Subject: [PATCH] v2: add defer enc.endEncoding() to all Metal encoders

---
 .../Embedding/EmbeddingGemmaModel.swift       | 22 +++++++++++++------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift b/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift
index 2467bfb..6a183b6 100644
--- a/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift
+++ b/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift
@@ -106,24 +106,32 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         guard !tokens.isEmpty else { return [] }
 
         let seqLen = tokens.count, hs = config.hiddenSize
+        print("  embed: seqLen=\(seqLen)")
 
         // Use single command buffer for entire forward pass
         let cmdBuf = engine.commandQueue.makeCommandBuffer()!
 
         // 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)") }
             hidden = try forwardLayer(hidden: hidden, layerIdx: idx, seqLen: seqLen, cmdBuf: cmdBuf)
         }
+        print("  embed: all layers done")
 
         // 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)
@@ -171,6 +179,7 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         let seqLen = tokens.count, hs = config.hiddenSize
         let buf = engine.device.makeBuffer(length: seqLen * hs * 4)!
         let enc = cmdBuf.makeComputeCommandEncoder()!
+        defer { enc.endEncoding() }
         let pso = try engine.pipeline(named: "lookup_embeddings")
         enc.setComputePipelineState(pso)
         enc.setBuffer(embedTokens, offset: 0, index: 0)
@@ -182,13 +191,13 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         enc.setBytes(&v, length: 4, index: 5)
         enc.dispatchThreads(MTLSize(width: seqLen, height: 1, depth: 1),
                            threadsPerThreadgroup: MTLSize(width: min(256, seqLen), height: 1, depth: 1))
-        enc.endEncoding()
         return buf
     }
 
     private func applyRmsNorm(input: MTLBuffer, weight: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
         let output = engine.device.makeBuffer(length: count * 4)!
         let enc = cmdBuf.makeComputeCommandEncoder()!
+        defer { enc.endEncoding() }
         let pso = try engine.pipeline(named: "rms_norm")
         enc.setComputePipelineState(pso)
         enc.setBuffer(input, offset: 0, index: 0)
@@ -199,12 +208,12 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         enc.setBytes(&e, length: 4, index: 4)
         enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
                            threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
-        enc.endEncoding()
         return output
     }
 
     private func matmulSeq(input: MTLBuffer, weight: MTLBuffer, output: MTLBuffer, m: Int, k: Int, n: Int, cmdBuf: MTLCommandBuffer) throws {
         let enc = cmdBuf.makeComputeCommandEncoder()!
+        defer { enc.endEncoding() }
         let pso = try engine.pipeline(named: "matmul_f32")
         enc.setComputePipelineState(pso)
         for i in 0..<m {
@@ -220,11 +229,11 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
             enc.dispatchThreads(MTLSize(width: n, height: 1, depth: 1),
                                threadsPerThreadgroup: MTLSize(width: min(256, n), height: 1, depth: 1))
         }
-        enc.endEncoding()
     }
 
     private func eltwiseAdd(a: MTLBuffer, b: MTLBuffer, output: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws {
         let enc = cmdBuf.makeComputeCommandEncoder()!
+        defer { enc.endEncoding() }
         let pso = try engine.pipeline(named: "eltwise_add")
         enc.setComputePipelineState(pso)
         enc.setBuffer(a, offset: 0, index: 0)
@@ -234,11 +243,11 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         enc.setBytes(&c, length: 4, index: 3)
         enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
                            threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
-        enc.endEncoding()
     }
 
     private func geluMul(gate: MTLBuffer, up: MTLBuffer, output: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws {
         let enc = cmdBuf.makeComputeCommandEncoder()!
+        defer { enc.endEncoding() }
         let pso = try engine.pipeline(named: "gelu_mul_kernel")
         enc.setComputePipelineState(pso)
         enc.setBuffer(gate, offset: 0, index: 0)
@@ -248,11 +257,11 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         enc.setBytes(&c, length: 4, index: 3)
         enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
                            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()!
+        defer { enc.endEncoding() }
         let pso = try engine.pipeline(named: "apply_rope")
         enc.setComputePipelineState(pso)
         enc.setBuffer(q, offset: 0, index: 0)
@@ -265,11 +274,11 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         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()!
+        defer { enc.endEncoding() }
         let pso = try engine.pipeline(named: "bidirectional_sliding_attn")
         enc.setComputePipelineState(pso)
         enc.setBuffer(q, offset: 0, index: 0)
@@ -289,7 +298,6 @@ public final class EmbeddingGemmaModel: @unchecked Sendable {
         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 {