v2: add EmbeddingGemma model with RoPE, sliding window attention, Q/K norm

Sources/MarkBase/Embedding/EmbeddingGemmaConfig.swift

@@ -1,34 +0,0 @@
-import Foundation
-
-/// EmbeddingGemma model configuration
-public struct EmbeddingGemmaConfig: Codable {
-    public let hiddenSize: Int
-    public let numHiddenLayers: Int
-    public let vocabSize: Int
-    public let numAttentionHeads: Int
-    public let numKeyValueHeads: Int
-    public let headDim: Int
-    public let intermediateSize: Int
-    public let maxPositionEmbeddings: Int
-    public let slidingWindow: Int
-    public let rmsNormEps: Float
-    public let ropeTheta: Float
-    public let useBidirectionalAttention: Bool
-    public let layerTypes: [String]
-
-    enum CodingKeys: String, CodingKey {
-        case hiddenSize = "hidden_size"
-        case numHiddenLayers = "num_hidden_layers"
-        case vocabSize = "vocab_size"
-        case numAttentionHeads = "num_attention_heads"
-        case numKeyValueHeads = "num_key_value_heads"
-        case headDim = "head_dim"
-        case intermediateSize = "intermediate_size"
-        case maxPositionEmbeddings = "max_position_embeddings"
-        case slidingWindow = "sliding_window"
-        case rmsNormEps = "rms_norm_eps"
-        case ropeTheta = "rope_theta"
-        case useBidirectionalAttention = "use_bidirectional_attention"
-        case layerTypes = "layer_types"
-    }
-}

Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift

@@ -0,0 +1,372 @@
+import Foundation
+import Metal
+import Accelerate
+
+/// EmbeddingGemmaConfig - Configuration for EmbeddingGemma model
+public struct EmbeddingGemmaConfig: Codable {
+    public let hiddenSize: Int
+    public let numHiddenLayers: Int
+    public let vocabSize: Int
+    public let numAttentionHeads: Int
+    public let numKeyValueHeads: Int
+    public let headDim: Int
+    public let intermediateSize: Int
+    public let maxPositionEmbeddings: Int
+    public let slidingWindow: Int
+    public let rmsNormEps: Float
+    public let ropeTheta: Float
+    public let useBidirectionalAttention: Bool
+    public let layerTypes: [String]
+
+    enum CodingKeys: String, CodingKey {
+        case hiddenSize = "hidden_size"
+        case numHiddenLayers = "num_hidden_layers"
+        case vocabSize = "vocab_size"
+        case numAttentionHeads = "num_attention_heads"
+        case numKeyValueHeads = "num_key_value_heads"
+        case headDim = "head_dim"
+        case intermediateSize = "intermediate_size"
+        case maxPositionEmbeddings = "max_position_embeddings"
+        case slidingWindow = "sliding_window"
+        case rmsNormEps = "rms_norm_eps"
+        case ropeTheta = "rope_theta"
+        case useBidirectionalAttention = "use_bidirectional_attention"
+        case layerTypes = "layer_types"
+    }
+
+    public static func load(from modelDir: String) throws -> Self {
+        let url = URL(fileURLWithPath: modelDir).appendingPathComponent("config.json")
+        let data = try Data(contentsOf: url)
+        return try JSONDecoder().decode(Self.self, from: data)
+    }
+}
+
+/// EmbeddingGemma - Google's 300M parameter embedding model
+public final class EmbeddingGemmaModel {
+    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]] = []
+    public var qProjs: [MTLBuffer] = []
+    public var kProjs: [MTLBuffer] = []
+    public var vProjs: [MTLBuffer] = []
+    public var oProjs: [MTLBuffer] = []
+    public var qNorms: [MTLBuffer] = []
+    public var kNorms: [MTLBuffer] = []
+    public var gateProjs: [MTLBuffer] = []
+    public var upProjs: [MTLBuffer] = []
+    public var downProjs: [MTLBuffer] = []
+
+    public init(modelDir: String, engine: MarkBaseEngine) throws {
+        self.engine = engine
+        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))")
+    }
+
+    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)!
+
+        for i in 0..<config.numHiddenLayers {
+            let p = "layers.\(i)"
+            layerNorms.append([
+                try loadBuffer("\(p).input_layernorm.weight"),
+                try loadBuffer("\(p).pre_feedforward_layernorm.weight"),
+                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]
+        }
+
+        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
+
+        // Embedding lookup
+        let inputBuf = try lookupEmbeddings(tokens: tokens)
+
+        // Forward through layers
+        var hidden = inputBuf
+        for idx in 0..<config.numHiddenLayers {
+            hidden = try forwardLayer(hidden: hidden, layerIdx: idx, seqLen: seqLen)
+        }
+
+        // Final norm
+        let output = try applyRmsNorm(input: hidden, weight: finalNorm, count: seqLen * hs)
+
+        // 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
+            for j in 0..<hs { embedding[j] += data[start + j] }
+        }
+        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
+
+    private func readTensor(_ name: String) throws -> [Float] {
+        guard let desc = reader.tensor(named: name) else {
+            throw WeightError.tensorNotFound(name)
+        }
+        let data = try reader.read(tensor: desc)
+        switch desc.dtype {
+        case .f32:
+            return data.withUnsafeBytes { Array(UnsafeBufferPointer(start: $0.baseAddress?.assumingMemoryBound(to: Float.self), count: data.count/4)) }
+        case .bf16:
+            return try SafeTensorsReader.bf16ToFloat32(data)
+        default:
+            throw WeightError.unsupportedDtype(desc.dtype.rawValue)
+        }
+    }
+
+    private func loadBuffer(_ name: String) throws -> MTLBuffer {
+        let data = try readTensor(name)
+        return engine.device.makeBuffer(bytes: data, length: data.count * 4)!
+    }
+
+    private func lookupEmbeddings(tokens: [Int]) 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)
+        enc.setBuffer(embedTokens, offset: 0, index: 0)
+        enc.setBytes(tokens, length: seqLen * MemoryLayout<Int>.size, index: 1)
+        enc.setBuffer(buf, offset: 0, index: 2)
+        var h = UInt32(hs), s = UInt32(seqLen), v = UInt32(config.vocabSize)
+        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),
+                           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 {
+        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)
+        enc.setBuffer(input, offset: 0, index: 0)
+        enc.setBuffer(weight, offset: 0, index: 1)
+        enc.setBuffer(output, offset: 0, index: 2)
+        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),
+                           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)
+        enc.setBuffer(input, offset: 0, index: 0)
+        enc.setBuffer(weight, offset: 0, index: 1)
+        enc.setBuffer(output, offset: 0, index: 2)
+        var mm = UInt32(m), kk = UInt32(k), nn = UInt32(n)
+        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: m * n, height: 1, depth: 1),
+                           threadsPerThreadgroup: MTLSize(width: min(256, m * n), height: 1, depth: 1))
+        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")
+        enc.setComputePipelineState(pso)
+        enc.setBuffer(a, offset: 0, index: 0)
+        enc.setBuffer(b, offset: 0, index: 1)
+        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),
+                           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()!
+        let pso = try engine.pipeline(named: "gelu_mul_kernel")
+        enc.setComputePipelineState(pso)
+        enc.setBuffer(gate, offset: 0, index: 0)
+        enc.setBuffer(up, offset: 0, index: 1)
+        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),
+                           threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
+        enc.endEncoding()
+    }
+}

Sources/MarkBase/Metal/EmbeddingKernels.metal

@@ -0,0 +1,189 @@
+#include <metal_stdlib>
+using namespace metal;
+
+// ── RoPE (Rotary Position Embedding) ──
+// Applies rotary position embeddings to Q and K tensors
+// Input: [seqLen, hiddenSize], positions: [seqLen]
+// Output: in-place modified Q/K
+
+kernel void apply_rope(
+    device float *q           [[buffer(0)]],
+    device float *k           [[buffer(1)]],
+    constant uint &seqLen     [[buffer(2)]],
+    constant uint &headDim    [[buffer(3)]],
+    constant uint &numHeads   [[buffer(4)]],
+    constant float &ropeTheta [[buffer(5)]],
+    uint tid [[thread_position_in_grid]],
+    uint gid [[threadgroup_position_in_grid]]
+) {
+    uint totalThreads = numHeads * headDim / 2;
+    if (tid >= totalThreads) return;
+
+    uint headIdx = tid / (headDim / 2);
+    uint dimIdx = tid % (headDim / 2);
+
+    // Each thread handles one (head, dim/2) pair across all positions
+    for (uint pos = 0; pos < seqLen; pos++) {
+        float theta = pow(ropeTheta, -2.0 * float(dimIdx) / float(headDim));
+        float freq = float(pos) * theta;
+        float cosFreq = cos(freq);
+        float sinFreq = sin(freq);
+
+        uint qBase = pos * numHeads * headDim + headIdx * headDim;
+        uint kBase = pos * numHeads * headDim + headIdx * headDim;
+
+        // Q rotation
+        float q0 = q[qBase + dimIdx];
+        float q1 = q[qBase + dimIdx + headDim / 2];
+        q[qBase + dimIdx] = q0 * cosFreq - q1 * sinFreq;
+        q[qBase + dimIdx + headDim / 2] = q0 * sinFreq + q1 * cosFreq;
+
+        // K rotation
+        float k0 = k[kBase + dimIdx];
+        float k1 = k[kBase + dimIdx + headDim / 2];
+        k[kBase + dimIdx] = k0 * cosFreq - k1 * sinFreq;
+        k[kBase + dimIdx + headDim / 2] = k0 * sinFreq + k1 * cosFreq;
+    }
+}
+
+// ── Q/K RMSNorm ──
+// Applies RMSNorm to each head's Q/K vector
+
+kernel void rms_norm_head(
+    device const float *input   [[buffer(0)]],
+    device const float *weight  [[buffer(1)]],
+    device float *output        [[buffer(2)]],
+    constant uint &headDim      [[buffer(3)]],
+    constant uint &numHeads     [[buffer(4)]],
+    constant float &eps         [[buffer(5)]],
+    uint tid [[thread_position_in_grid]]
+) {
+    if (tid >= numHeads) return;
+
+    uint base = tid * headDim;
+    float sumSq = 0.0;
+    for (uint i = 0; i < headDim; i++) {
+        sumSq += input[base + i] * input[base + i];
+    }
+    float rms = sqrt(sumSq / float(headDim) + eps);
+
+    for (uint i = 0; i < headDim; i++) {
+        output[base + i] = input[base + i] * weight[i] / rms;
+    }
+}
+
+// ── Bidirectional Sliding Window Attention ──
+// Computes softmax(Q*K^T/sqrt(d)) * V with sliding window mask
+
+kernel void bidirectional_sliding_attn(
+    device const float *q           [[buffer(0)]],
+    device const float *k           [[buffer(1)]],
+    device const float *v           [[buffer(2)]],
+    device float *output            [[buffer(3)]],
+    constant uint &seqLen           [[buffer(4)]],
+    constant uint &headDim          [[buffer(5)]],
+    constant uint &numHeads         [[buffer(6)]],
+    constant uint &numKVHeads       [[buffer(7)]],
+    constant uint &slidingWindow    [[buffer(8)]],
+    constant float &scale           [[buffer(9)]],
+    threadgroup float *shared_mem   [[threadgroup(0)]],
+    uint tid [[thread_position_in_grid]],
+    uint tgSize [[threads_per_threadgroup]]
+) {
+    // Each thread handles one (query_position, head) pair
+    uint totalQueries = seqLen * numHeads;
+    if (tid >= totalQueries) return;
+
+    uint qPos = tid / numHeads;
+    uint headIdx = tid % numHeads;
+    uint kvHeadIdx = headIdx * numKVHeads / numHeads;  // GQA
+
+    float sqrtD = sqrt(float(headDim));
+    uint kvBase = kvHeadIdx * headDim;
+
+    // Compute attention scores with sliding window
+    float maxScore = -1e30f;
+    float scores[2048];  // max seqLen
+    uint validCount = 0;
+
+    uint windowStart = qPos > slidingWindow ? qPos - slidingWindow : 0;
+    uint windowEnd = min(qPos + slidingWindow + 1, seqLen);
+
+    for (uint kPos = windowStart; kPos < windowEnd; kPos++) {
+        float dot = 0.0;
+        for (uint d = 0; d < headDim; d++) {
+            dot += q[qPos * numHeads * headDim + headIdx * headDim + d] *
+                   k[kPos * numKVHeads * headDim + kvBase + d];
+        }
+        scores[validCount] = dot * scale / sqrtD;
+        if (scores[validCount] > maxScore) maxScore = scores[validCount];
+        validCount++;
+    }
+
+    // Softmax
+    float sumExp = 0.0;
+    for (uint i = 0; i < validCount; i++) {
+        scores[i] = exp(scores[i] - maxScore);
+        sumExp += scores[i];
+    }
+    if (sumExp > 0) {
+        for (uint i = 0; i < validCount; i++) {
+            scores[i] /= sumExp;
+        }
+    }
+
+    // Weighted sum of V
+    uint outBase = qPos * numHeads * headDim + headIdx * headDim;
+    for (uint d = 0; d < headDim; d++) {
+        float val = 0.0;
+        uint kPosIdx = windowStart;
+        for (uint i = 0; i < validCount; i++) {
+            val += scores[i] * v[kPosIdx * numKVHeads * headDim + kvBase + d];
+            kPosIdx++;
+        }
+        output[outBase + d] = val;
+    }
+}
+
+// ── GELU ──
+
+kernel void gelu_kernel(
+    device const float *input  [[buffer(0)]],
+    device float *output       [[buffer(1)]],
+    constant uint &count       [[buffer(2)]],
+    uint tid [[thread_position_in_grid]]
+) {
+    if (tid >= count) return;
+    float x = input[tid];
+    float absv = abs(x);
+    float gelu;
+    if (absv > 10.0f) {
+        gelu = x > 0 ? x : 0.0f;
+    } else {
+        float x3 = x * x * x;
+        gelu = 0.5f * x * (1.0f + tanh(0.7978845608028654f * (x + 0.044715f * x3)));
+    }
+    output[tid] = gelu;
+}
+
+// ── GELU + Multiply ──
+
+kernel void gelu_mul_kernel(
+    device const float *gate   [[buffer(0)]],
+    device const float *up     [[buffer(1)]],
+    device float *output       [[buffer(2)]],
+    constant uint &count       [[buffer(3)]],
+    uint tid [[thread_position_in_grid]]
+) {
+    if (tid >= count) return;
+    float g = gate[tid];
+    float absv = abs(g);
+    float gelu;
+    if (absv > 10.0f) {
+        gelu = g > 0 ? g : 0.0f;
+    } else {
+        float g3 = g * g * g;
+        gelu = 0.5f * g * (1.0f + tanh(0.7978845608028654f * (g + 0.044715f * g3)));
+    }
+    output[tid] = gelu * up[tid];
+}