commit 8a66b9086af0dcf17c169be627b209669886d018 Author: MarkBase Admin Date: Sun Jul 5 13:29:25 2026 +0800 v2: Initial clean branch with unit tests + CI/CD pipeline - Started from ac75faa (initial E4B-MarkBase integration) - Kept Sources/ (all engine code) + Package.swift + .gitignore - Removed all ad-hoc tests, documentation, scripts, Python files - Added Tests/00_Unit/ (MathTest, TokenizerTest, SamplerTest) - Added .gitea/workflows/ci.yaml (build + unit tests + lint) - Added Scripts/check_resources.sh (memory-aware test runner) - Added Tests/Manifest.json (resource requirements for all tests) - Focus: 4-bit quantized models only diff --git a/.gitea/workflows/ci.yaml b/.gitea/workflows/ci.yaml new file mode 100644 index 0000000..4159735 --- /dev/null +++ b/.gitea/workflows/ci.yaml @@ -0,0 +1,42 @@ +name: CI + +on: + push: + branches: [ v2 ] + pull_request: + branches: [ v2 ] + +jobs: + build: + runs-on: macos-latest + steps: + - uses: actions/checkout@v4 + + - name: Build Swift + run: swift build -c debug + + - name: Build Release + run: swift build -c release + + unit-tests: + needs: build + runs-on: macos-latest + steps: + - uses: actions/checkout@v4 + + - name: Run Unit Tests + run: swift test --filter "00_Unit" + + lint: + needs: build + runs-on: macos-latest + steps: + - uses: actions/checkout@v4 + + - name: Check for debug prints + run: | + if grep -r "print(" Sources/MarkBase/ --include="*.swift" | grep -v "//.*print" | grep -v "Error"; then + echo "WARNING: Debug print() found in Sources/" + exit 0 + fi + echo "No debug prints found" diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..0d66706 --- /dev/null +++ b/.gitignore @@ -0,0 +1,10 @@ +.build/ +models/ +*.log +DerivedData/ +.swiftpm/ +Package.resolved +*.xcodeproj/ +*.xcworkspace/ +.DS_Store +test_summary.md \ No newline at end of file diff --git a/Package.swift b/Package.swift new file mode 100644 index 0000000..1c34907 --- /dev/null +++ b/Package.swift @@ -0,0 +1,50 @@ +// swift-tools-version: 6.0 +import PackageDescription + +let package = Package( + name: "MarkBase", + platforms: [.macOS(.v15)], + products: [ + .library(name: "MarkBase", targets: ["MarkBase"]), + .executable(name: "MarkBaseServer", targets: ["MarkBaseServer"]), + .executable(name: "CLITest", targets: ["CLITest"]), + ], + dependencies: [ + .package(url: "https://github.com/hummingbird-project/hummingbird.git", from: "2.0.0"), + .package(path: "/Users/accusys/coder/poc/rdma"), + ], + targets: [ + .target( + name: "MarkBase", + exclude: ["Metal/MetalKernels.metal", "Metal/OptimizedKernels.metal", "Metal/FusionKernels.metal", "Metal/MetalKernels.metallib", "Metal/metallib"], + linkerSettings: [ + .linkedFramework("Metal"), + .linkedFramework("Foundation"), + ] + ), + .executableTarget( + name: "MarkBaseServer", + dependencies: [ + "MarkBase", + .product(name: "Hummingbird", package: "hummingbird"), + .product(name: "RDMAKit", package: "rdma"), + ], + linkerSettings: [ + .linkedFramework("Metal"), + .linkedFramework("Foundation"), + ] + ), + .executableTarget( + name: "CLITest", + dependencies: ["MarkBase"], + linkerSettings: [ + .linkedFramework("Metal"), + .linkedFramework("Foundation"), + ] + ), + .testTarget( + name: "MarkBaseTests", + dependencies: ["MarkBase"] + ), + ] +) diff --git a/Scripts/check_resources.sh b/Scripts/check_resources.sh new file mode 100755 index 0000000..4462126 --- /dev/null +++ b/Scripts/check_resources.sh @@ -0,0 +1,30 @@ +#!/bin/bash +# check_resources.sh — Check available system memory before running tests +# Usage: check_resources.sh + +set -euo pipefail + +REQUIRED_GB="${1:-0}" + +# Get available memory using vm_stat on macOS +if [[ "$(uname)" == "Darwin" ]]; then + PAGE_SIZE=$(vm_stat | head -1 | awk '{print $NF}' | sed 's/\.//') + FREE_PAGES=$(vm_stat | awk '/free/ {print $NF}' | sed 's/\.//') + if [[ -z "$FREE_PAGES" ]]; then + FREE_PAGES=$(vm_stat | awk '/Pages free/ {print $3}' | sed 's/\.//') + fi + AVAILABLE_GB=$(( FREE_PAGES * 16384 / 1073741824 )) # 16384 = page size on Apple Silicon + echo "Available memory: ~${AVAILABLE_GB}GB (free: ${FREE_PAGES} pages)" +else + AVAILABLE_GB=$(free -g | awk '/^Mem:/ {print $7}') + echo "Available memory: ~${AVAILABLE_GB}GB" +fi + +if [[ "$AVAILABLE_GB" -lt "$REQUIRED_GB" ]]; then + echo "ERROR: Need ${REQUIRED_GB}GB but only ${AVAILABLE_GB}GB available" + echo "Run 'memory_pressure' to check memory status" + exit 1 +fi + +echo "OK: ${AVAILABLE_GB}GB >= ${REQUIRED_GB}GB required" +exit 0 diff --git a/Sources/CLITest/main.swift b/Sources/CLITest/main.swift new file mode 100644 index 0000000..d0b7ecf --- /dev/null +++ b/Sources/CLITest/main.swift @@ -0,0 +1,52 @@ +import Foundation +import MarkBase + +// Simple CLI test for forward pass +let modelDir = "./models/gemma-4-12b-it-4bit" + +guard FileManager.default.fileExists(atPath: modelDir + "/config.json") else { + print("Model not found at \(modelDir)") + exit(1) +} + +print("Loading engine...") +let engine = try MarkBaseEngine(autoCompile: true) +print("✓ Engine created") + +print("\nLoading 12B model...") +let start = Date() +let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 128) +let loadTime = Date().timeIntervalSince(start) + +print("✓ Model loaded in \(String(format: "%.1f", loadTime))s") +print(" Layers: \(model.numHiddenLayers)") +print(" Hidden: \(model.hiddenSize)") +print(" Vocab: \(model.vocabSize)") + +// Test forward pass +print("\n=== Testing forward pass ===") +print("Testing token 2 (BOS) at position 0...") +let logits = try model.forward(tokenId: 2, position: 0, debug: true) + +print("\n✓ Forward pass complete: \(logits.count) logits") + +let maxLogit = logits.max() ?? -999 +let minLogit = logits.min() ?? -999 +let hasNaN = logits.contains { $0.isNaN } +let nanCount = logits.filter { $0.isNaN }.count + +print(" Max logit: \(maxLogit)") +print(" Min logit: \(minLogit)") +print(" NaN count: \(nanCount)/\(logits.count)") +print(" Has NaN: \(hasNaN)") + +if !hasNaN { + let sorted = logits.enumerated().sorted { $0.element > $1.element } + let top10 = sorted.prefix(10) + print("\n Top 10 tokens:") + for (idx, logit) in top10 { + print(" Token \(idx): \(String(format: "%.4f", logit))") + } +} + +print("\n✅ CLI test complete!") \ No newline at end of file diff --git a/Sources/MarkBase/AsyncInference.swift b/Sources/MarkBase/AsyncInference.swift new file mode 100644 index 0000000..3be54d2 --- /dev/null +++ b/Sources/MarkBase/AsyncInference.swift @@ -0,0 +1,190 @@ +import Foundation +import Metal + +// ───────────────────────────────────────────────────────────── +// Async Inference Optimizations +// ───────────────────────────────────────────────────────────── + +/// Async inference result +public struct InferenceResult { + public let logits: [Float] + public let elapsed: TimeInterval + public let token: Int + + public init(logits: [Float], elapsed: TimeInterval, token: Int = 0) { + self.logits = logits + self.elapsed = elapsed + self.token = token + } +} + +/// Async inference queue for batching requests +public final class AsyncInferenceQueue { + private let maxBatchSize: Int + private var pending: [(input: Int, position: Int, completion: (Result) -> Void)] = [] + private let lock = NSLock() + + public init(maxBatchSize: Int = 8) { + self.maxBatchSize = maxBatchSize + } + + /// Add inference request + public func enqueue(input: Int, position: Int, completion: @escaping (Result) -> Void) { + lock.lock() + pending.append((input, position, completion)) + let shouldProcess = pending.count >= maxBatchSize + lock.unlock() + + if shouldProcess { + processBatch() + } + } + + /// Process batch of requests + private func processBatch() { + lock.lock() + let batch = pending.prefix(maxBatchSize) + pending.removeFirst(min(batch.count, maxBatchSize)) + lock.unlock() + + // TODO: Implement actual batch processing + // This would require modifications to the model to support batch inference + } +} + +/// Async token generator with prefetching +public final class AsyncTokenGenerator: @unchecked Sendable { + private let model: E4BModel + private let tokenizer: Tokenizer + private let engine: MarkBaseEngine + private let sampler: Sampler + + private var currentLogits: [Float]? + private var currentPosition: Int = 0 + private var generatedTokens: [Int] = [] + + public init(model: E4BModel, tokenizer: Tokenizer, engine: MarkBaseEngine) { + self.model = model + self.tokenizer = tokenizer + self.engine = engine + self.sampler = Sampler() + } + + /// Start generation with async streaming + public func start(prompt: String, config: GenerationConfig) -> AsyncStream { + return AsyncStream { [weak self] continuation in + guard let self = self else { + continuation.finish() + return + } + + Task.detached { + do { + try await self.generateAsync(prompt: prompt, config: config) { tokenText in + continuation.yield(tokenText) + } + continuation.finish() + } catch { + continuation.finish() + } + } + } + } + + /// Async generation with callback + private func generateAsync( + prompt: String, + config: GenerationConfig, + onToken: (String) -> Void + ) async throws { + let promptTokens = tokenizer.encode(text: prompt) + + // Pre-fill KV cache + var lastLogits: [Float] = [] + for (position, tokenId) in promptTokens.enumerated() { + lastLogits = try model.forward(tokenId: tokenId, position: position) + } + + currentLogits = lastLogits + currentPosition = promptTokens.count + generatedTokens = [] + var streamDecoder = StreamingDecoder(tokenizer: tokenizer) + + // Generate tokens + for _ in 0.. [Float] { + if let logits = prefetchedLogits, nextToken == token, nextPosition == position { + prefetchedLogits = nil + return logits + } + + // Compute on demand + return try model.forward(tokenId: token, position: position) + } +} + diff --git a/Sources/MarkBase/Audio/AudioConfig.swift b/Sources/MarkBase/Audio/AudioConfig.swift new file mode 100644 index 0000000..abb1dc2 --- /dev/null +++ b/Sources/MarkBase/Audio/AudioConfig.swift @@ -0,0 +1,51 @@ +import Foundation + +public struct AudioConfig: Codable { + public let hiddenSize: Int + public let numAttentionHeads: Int + public let numHiddenLayers: Int + public let convKernelSize: Int + public let attentionChunkSize: Int + public let attentionContextLeft: Int + public let attentionContextRight: Int + public let attentionLogitCap: Float + public let hiddenAct: String + public let rmsNormEps: Float + public let outputProjDims: Int + public let subsamplingConvChannels: [Int] + public let residualWeight: Float + + public init( + hiddenSize: Int = 1024, + numAttentionHeads: Int = 8, + numHiddenLayers: Int = 12, + convKernelSize: Int = 5, + attentionChunkSize: Int = 12, + attentionContextLeft: Int = 13, + attentionContextRight: Int = 0, + attentionLogitCap: Float = 50.0, + hiddenAct: String = "silu", + rmsNormEps: Float = 1e-6, + outputProjDims: Int = 1536, + subsamplingConvChannels: [Int] = [128, 32], + residualWeight: Float = 0.5 + ) { + self.hiddenSize = hiddenSize + self.numAttentionHeads = numAttentionHeads + self.numHiddenLayers = numHiddenLayers + self.convKernelSize = convKernelSize + self.attentionChunkSize = attentionChunkSize + self.attentionContextLeft = attentionContextLeft + self.attentionContextRight = attentionContextRight + self.attentionLogitCap = attentionLogitCap + self.hiddenAct = hiddenAct + self.rmsNormEps = rmsNormEps + self.outputProjDims = outputProjDims + self.subsamplingConvChannels = subsamplingConvChannels + self.residualWeight = residualWeight + } + + public var headDim: Int { + hiddenSize / numAttentionHeads + } +} \ No newline at end of file diff --git a/Sources/MarkBase/Audio/AudioFeatureExtractor.swift b/Sources/MarkBase/Audio/AudioFeatureExtractor.swift new file mode 100644 index 0000000..82f862d --- /dev/null +++ b/Sources/MarkBase/Audio/AudioFeatureExtractor.swift @@ -0,0 +1,230 @@ +import AVFoundation +import Metal + +public final class AudioFeatureExtractor { + public let sampleRate: Int + public let nMels: Int + public let nFft: Int + public let hopLength: Int + public let fMin: Float + public let fMax: Float + + public init( + sampleRate: Int = 16000, + nMels: Int = 128, + nFft: Int = 400, + hopLength: Int = 160, + fMin: Float = 0, + fMax: Float = 8000 + ) { + self.sampleRate = sampleRate + self.nMels = nMels + self.nFft = nFft + self.hopLength = hopLength + self.fMin = fMin + self.fMax = fMax + } + + public func extractMelSpectrogram(from audioData: [Float]) -> [[Float]] { + let numFrames = (audioData.count - nFft) / hopLength + 1 + var melSpec = [[Float]](repeating: [Float](repeating: 0, count: nMels), count: numFrames) + + for frameIdx in 0.. [Float] { + let n = frame.count + return frame.enumerated().map { i, val in + val * 0.5 * (1.0 - cos(2.0 * Float.pi * Float(i) / Float(n - 1))) + } + } + + private func computeSpectrum(_ frame: [Float]) -> [Float] { + let n = frame.count + var spectrum = [Float](repeating: 0, count: n / 2 + 1) + + for k in 0.. [Float] { + var melEnergies = [Float](repeating: 0, count: nMels) + + let melPoints = createMelFilterbank() + + for melIdx in 0.. [[Float]] { + var filterbank = [[Float]](repeating: [Float](repeating: 0, count: nFft / 2 + 1), count: nMels) + + let melMin = hzToMel(fMin) + let melMax = hzToMel(fMax) + + let melPoints = (0.. Float { + 2595.0 * log10(1.0 + hz / 700.0) + } + + private func melToHz(_ mel: Float) -> Float { + 700.0 * (pow(10.0, mel / 2595.0) - 1.0) + } + + // ── GPU-accelerated mel spectrogram ── + + public func extractMelSpectrogramGPU( + engine: MarkBaseEngine, + audioData: [Float] + ) throws -> [[Float]] { + let device = engine.device + let spectrumSize = nFft / 2 + 1 + let numFrames = (audioData.count - nFft) / hopLength + 1 + let melBufferSize = numFrames * nMels + + let filterbank2D = createMelFilterbank() + var flatFilterbank = [Float](repeating: 0, count: nMels * spectrumSize) + for m in 0.. [Float] { + let asset = AVURLAsset(url: url) + let reader = try AVAssetReader(asset: asset) + + let output = AVAssetReaderAudioMixOutput(audioTracks: asset.tracks, audioSettings: nil) + reader.add(output) + reader.startReading() + + var samples: [Float] = [] + while reader.status == .reading { + let buffer = output.copyNextSampleBuffer() + if let buffer = buffer { + let blockBuffer = CMSampleBufferGetDataBuffer(buffer) + if let blockBuffer = blockBuffer { + let length = CMBlockBufferGetDataLength(blockBuffer) + var data = [Float](repeating: 0, count: length / MemoryLayout.stride) + CMBlockBufferCopyDataBytes(blockBuffer, atOffset: 0, dataLength: length, destination: &data) + samples.append(contentsOf: data) + } + } + } + + return samples + } +} \ No newline at end of file diff --git a/Sources/MarkBase/Audio/AudioTower.swift b/Sources/MarkBase/Audio/AudioTower.swift new file mode 100644 index 0000000..659afa7 --- /dev/null +++ b/Sources/MarkBase/Audio/AudioTower.swift @@ -0,0 +1,740 @@ +import Metal + +public final class AudioTower { + public let config: AudioConfig + public let engine: MarkBaseEngine + public let weights: AudioWeights + + private var normBuffer: MTLBuffer + private var qBuffer: MTLBuffer + private var kBuffer: MTLBuffer + private var vBuffer: MTLBuffer + private var attnOutBuffer: MTLBuffer + private var ffnBuffer: MTLBuffer + private var tempBuffer: MTLBuffer + private var subsampleBuf: MTLBuffer + private var layerBuffer: MTLBuffer // NEW: dedicated buffer for audio layers + + public init(config: AudioConfig, engine: MarkBaseEngine, weights: AudioWeights) throws { + self.config = config + self.engine = engine + self.weights = weights + + let device = engine.device + let maxSeqLen = 4096 + let hiddenSize = config.hiddenSize + let headDim = config.headDim + + normBuffer = device.makeBuffer(length: hiddenSize * maxSeqLen * 4)! + qBuffer = device.makeBuffer(length: hiddenSize * maxSeqLen * 4)! + kBuffer = device.makeBuffer(length: hiddenSize * maxSeqLen * 4)! + vBuffer = device.makeBuffer(length: hiddenSize * maxSeqLen * 4)! + attnOutBuffer = device.makeBuffer(length: hiddenSize * maxSeqLen * 4)! + ffnBuffer = device.makeBuffer(length: 4096 * maxSeqLen * 4)! + tempBuffer = device.makeBuffer(length: max(hiddenSize, 4096) * maxSeqLen * 4)! + subsampleBuf = device.makeBuffer(length: max(hiddenSize, 128 * 64) * maxSeqLen * 4)! + layerBuffer = device.makeBuffer(length: max(hiddenSize, 4096) * maxSeqLen * 4)! // NEW + } + + public func forward(inputBuffer: MTLBuffer, seqLen: Int, outputBuffer: MTLBuffer) throws { + var current = inputBuffer + var currentLen = seqLen + + let cmdBuf = engine.commandQueue.makeCommandBuffer()! + + // 1. Subsample conv: mel [seqLen, 128] -> [seqLen/4, 1024] + let (projInput, projLen) = try applySubsampleConv( + melInput: current, + nMels: 128, + seqLen: currentLen, + cmdBuf: cmdBuf + ) + + let cmdBuf2 = engine.commandQueue.makeCommandBuffer()! + + // 2. Input projection: [seqLen/4, 1024] -> [seqLen/4, 1024] + current = try applyInputProjection(input: projInput, seqLen: projLen, cmdBuf: cmdBuf2) + currentLen = projLen + + let cmdBuf3 = engine.commandQueue.makeCommandBuffer()! + + // 3. Audio layers (12 layers) + for layerWeights in weights.layers { + current = try applyLayer( + input: current, + weights: layerWeights, + seqLen: currentLen, + cmdBuf: cmdBuf3 + ) + } + + let cmdBuf4 = engine.commandQueue.makeCommandBuffer()! + + // 4. Output projection: [seqLen/4, 1024] -> [seqLen/4, 1536] + try applyOutputProjection(input: current, seqLen: currentLen, output: outputBuffer, cmdBuf: cmdBuf4) + + cmdBuf4.commit() + cmdBuf4.waitUntilCompleted() + } + + private func applySubsampleConv( + melInput: MTLBuffer, + nMels: Int, + seqLen: Int, + cmdBuf: MTLCommandBuffer + ) throws -> (MTLBuffer, Int) { + // Input mel: [seqLen, 128] row-major + // Step 1: Transpose to CHW [1, 128, seqLen] + let chwInput = try transposeMelToCHW(input: melInput, nMels: nMels, seqLen: seqLen, cmdBuf: cmdBuf) + + // Step 2: Layer0 conv2d [1, 128, seqLen] -> [128, 64, seqLen/2] + let layer0Out = try applyConv2DLayer( + input: chwInput, + inCh: 1, + height: nMels, + width: seqLen, + convWeight: weights.subsampleConvLayer0.convWeight, + normWeight: weights.subsampleConvLayer0.normWeight, + outChannels: 128, + outputBuffer: tempBuffer, + cmdBuf: cmdBuf + ) + let h1 = (nMels + 1) / 2 + let w1 = (seqLen + 1) / 2 + + // Step 3: Layer1 conv2d [128, 64, seqLen/2] -> [32, 32, seqLen/4] + let layer1Out = try applyConv2DLayer( + input: layer0Out, + inCh: 128, + height: h1, + width: w1, + convWeight: weights.subsampleConvLayer1.convWeight, + normWeight: weights.subsampleConvLayer1.normWeight, + outChannels: 32, + outputBuffer: subsampleBuf, + cmdBuf: cmdBuf + ) + let h2 = (h1 + 1) / 2 + let w2 = (w1 + 1) / 2 + + // Step 4: Flatten [32, 32, w2] -> [w2, 1024] + let flatOutput = try flattenCHW(input: layer1Out, C: 32, H: h2, W: w2, outputBuffer: tempBuffer, cmdBuf: cmdBuf) + + return (flatOutput, w2) + } + + private func transposeMelToCHW( + input: MTLBuffer, + nMels: Int, + seqLen: Int, + cmdBuf: MTLCommandBuffer + ) throws -> MTLBuffer { + let output = subsampleBuf + + let pso = try engine.pipeline(named: "transpose_2d") + let enc = cmdBuf.makeComputeCommandEncoder()! + enc.setComputePipelineState(pso) + + enc.setBuffer(input, offset: 0, index: 0) + enc.setBuffer(output, offset: 0, index: 1) + + // FIX: Input is [seqLen, nMels], transpose to [nMels, seqLen] + var rows = UInt32(seqLen) // FIX: was nMels, should be seqLen + enc.setBytes(&rows, length: 4, index: 2) + var cols = UInt32(nMels) // FIX: was seqLen, should be nMels + enc.setBytes(&cols, length: 4, index: 3) + + let grid = MTLSize(width: nMels, height: seqLen, depth: 1) // FIX: grid for output [nMels, seqLen] + let tg = engine.threadgroupSize2D(pso, grid: (nMels, seqLen)) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + + return output + } + + private func applyConv2DLayer( + input: MTLBuffer, + inCh: Int, + height: Int, + width: Int, + convWeight: MTLBuffer, + normWeight: MTLBuffer, + outChannels: Int, + outputBuffer: MTLBuffer, + cmdBuf: MTLCommandBuffer + ) throws -> MTLBuffer { + let pso = try engine.pipeline(named: "audio_subsample_conv_2d") + let enc = cmdBuf.makeComputeCommandEncoder()! + enc.setComputePipelineState(pso) + + enc.setBuffer(input, offset: 0, index: 0) + enc.setBuffer(convWeight, offset: 0, index: 1) + enc.setBuffer(normWeight, offset: 0, index: 2) + enc.setBuffer(outputBuffer, offset: 0, index: 3) + + var inCh_ = UInt32(inCh) + enc.setBytes(&inCh_, length: 4, index: 4) + var outCh_ = UInt32(outChannels) + enc.setBytes(&outCh_, length: 4, index: 5) + var h_ = UInt32(height) + enc.setBytes(&h_, length: 4, index: 6) + var w_ = UInt32(width) + enc.setBytes(&w_, length: 4, index: 7) + + let outH = (height + 1) / 2 + let outW = (width + 1) / 2 + let grid = MTLSize(width: outChannels, height: outH, depth: outW) + let tg = MTLSize(width: 8, height: 8, depth: 4) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + + return outputBuffer + } + + private func flattenCHW( + input: MTLBuffer, + C: Int, + H: Int, + W: Int, + outputBuffer: MTLBuffer, + cmdBuf: MTLCommandBuffer + ) throws -> MTLBuffer { + let pso = try engine.pipeline(named: "audio_flatten_chw") + let enc = cmdBuf.makeComputeCommandEncoder()! + enc.setComputePipelineState(pso) + + enc.setBuffer(input, offset: 0, index: 0) + enc.setBuffer(outputBuffer, offset: 0, index: 1) + + var C_ = UInt32(C) + enc.setBytes(&C_, length: 4, index: 2) + var H_ = UInt32(H) + enc.setBytes(&H_, length: 4, index: 3) + var W_ = UInt32(W) + enc.setBytes(&W_, length: 4, index: 4) + + let grid = MTLSize(width: C * H, height: W, depth: 1) + let tg = engine.threadgroupSize2D(pso, grid: (C * H, W)) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + + return outputBuffer + } + + private func applyInputProjection(input: MTLBuffer, seqLen: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer { + // FIX: Use subsampleBuf as output to avoid overwriting input (tempBuffer) + let output = subsampleBuf + + // Input: [seqLen, 1024] after flatten (32 channels * 32 height = 1024) + // Weight: [1024, 1024] float32 + // Output: [seqLen, 1024] (hiddenSize) + + let pso = try engine.pipeline(named: "audio_linear_seq") + let enc = cmdBuf.makeComputeCommandEncoder()! + enc.setComputePipelineState(pso) + + enc.setBuffer(input, offset: 0, index: 0) + enc.setBuffer(weights.inputProjLinearWeight, offset: 0, index: 1) + enc.setBuffer(nil, offset: 0, index: 2) // No bias + enc.setBuffer(output, offset: 0, index: 3) + + var inFeatures = UInt32(1024) + enc.setBytes(&inFeatures, length: 4, index: 4) + var outFeatures = UInt32(1024) + enc.setBytes(&outFeatures, length: 4, index: 5) + var hasBias = false + enc.setBytes(&hasBias, length: 1, index: 6) + var seqLen_ = UInt32(seqLen) + enc.setBytes(&seqLen_, length: 4, index: 7) + + let grid = MTLSize(width: 1024, height: seqLen, depth: 1) + let tg = engine.threadgroupSize2D(pso, grid: (1024, seqLen)) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + + return output + } + + private func applyLayer( + input: MTLBuffer, + weights: AudioLayerWeights, + seqLen: Int, + cmdBuf: MTLCommandBuffer + ) throws -> MTLBuffer { + var current = input + + // 1. Norm pre-attn + current = try applyRMSNorm( + input: current, + weight: weights.normPreAttn, + seqLen: seqLen, + hiddenSize: config.hiddenSize, + cmdBuf: cmdBuf + ) + + // 2. Self-attention with relative position + let attnOut = try applySelfAttention( + input: current, + weights: weights, + seqLen: seqLen, + cmdBuf: cmdBuf + ) + + // 3. Residual + norm post-attn + current = try applyResidualAdd( + input: input, + add: attnOut, + seqLen: seqLen, + hiddenSize: config.hiddenSize, + residualWeight: config.residualWeight, + cmdBuf: cmdBuf + ) + + current = try applyRMSNorm( + input: current, + weight: weights.normPostAttn, + seqLen: seqLen, + hiddenSize: config.hiddenSize, + cmdBuf: cmdBuf + ) + + // 4. Local conv1d + let lconvOut = try applyLConv1D( + input: current, + weights: weights, + seqLen: seqLen, + cmdBuf: cmdBuf + ) + + // 5. Residual + current = try applyResidualAdd( + input: current, + add: lconvOut, + seqLen: seqLen, + hiddenSize: config.hiddenSize, + residualWeight: config.residualWeight, + cmdBuf: cmdBuf + ) + + // 6. Feed-forward 1 + let ff1Out = try applyFeedForward( + input: current, + weights: weights.feedForward1, + seqLen: seqLen, + cmdBuf: cmdBuf + ) + + // 7. Residual + current = try applyResidualAdd( + input: current, + add: ff1Out, + seqLen: seqLen, + hiddenSize: config.hiddenSize, + residualWeight: config.residualWeight, + cmdBuf: cmdBuf + ) + + // 8. Feed-forward 2 + let ff2Out = try applyFeedForward( + input: current, + weights: weights.feedForward2, + seqLen: seqLen, + cmdBuf: cmdBuf + ) + + // 9. Residual + norm out + current = try applyResidualAdd( + input: current, + add: ff2Out, + seqLen: seqLen, + hiddenSize: config.hiddenSize, + residualWeight: config.residualWeight, + cmdBuf: cmdBuf + ) + + current = try applyRMSNorm( + input: current, + weight: weights.normOut, + seqLen: seqLen, + hiddenSize: config.hiddenSize, + cmdBuf: cmdBuf + ) + + return current + } + + private func applySelfAttention( + input: MTLBuffer, + weights: AudioLayerWeights, + seqLen: Int, + cmdBuf: MTLCommandBuffer + ) throws -> MTLBuffer { + // Q, K, V projections + let q = try applyQuantizedLinear( + input: input, + weights: weights.selfAttnQProj, + seqLen: seqLen, + output: qBuffer, + cmdBuf: cmdBuf + ) + + let k = try applyQuantizedLinear( + input: input, + weights: weights.selfAttnKProj, + seqLen: seqLen, + output: kBuffer, + cmdBuf: cmdBuf + ) + + let v = try applyQuantizedLinear( + input: input, + weights: weights.selfAttnVProj, + seqLen: seqLen, + output: vBuffer, + cmdBuf: cmdBuf + ) + + // Attention with relative position and context + let attnOut = try applyAudioAttention( + q: q, + k: k, + v: v, + relativeKProj: weights.selfAttnRelativeKProj, + perDimScale: weights.selfAttnPerDimScale, + seqLen: seqLen, + numHeads: config.numAttentionHeads, + headDim: config.headDim, + contextLeft: config.attentionContextLeft, + logitCap: config.attentionLogitCap, + output: attnOutBuffer, + cmdBuf: cmdBuf + ) + + // Post projection + let output = try applyQuantizedLinear( + input: attnOut, + weights: weights.selfAttnPost, + seqLen: seqLen, + output: tempBuffer, + cmdBuf: cmdBuf + ) + + return output + } + + private func applyAudioAttention( + q: MTLBuffer, + k: MTLBuffer, + v: MTLBuffer, + relativeKProj: MTLBuffer, + perDimScale: MTLBuffer, + seqLen: Int, + numHeads: Int, + headDim: Int, + contextLeft: Int, + logitCap: Float, + output: MTLBuffer, + cmdBuf: MTLCommandBuffer + ) throws -> MTLBuffer { + let pso = try engine.pipeline(named: "audio_attention_full") + let enc = cmdBuf.makeComputeCommandEncoder()! + 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(relativeKProj, offset: 0, index: 3) + enc.setBuffer(perDimScale, offset: 0, index: 4) + enc.setBuffer(output, offset: 0, index: 5) + + var seqLen_ = UInt32(seqLen) + enc.setBytes(&seqLen_, length: 4, index: 6) + var numHeads_ = UInt32(numHeads) + enc.setBytes(&numHeads_, length: 4, index: 7) + var headDim_ = UInt32(headDim) + enc.setBytes(&headDim_, length: 4, index: 8) + var contextLeft_ = UInt32(contextLeft) + enc.setBytes(&contextLeft_, length: 4, index: 9) + var logitCap_ = logitCap + enc.setBytes(&logitCap_, length: 4, index: 10) + + let grid = MTLSize(width: numHeads * headDim, height: seqLen, depth: 1) + let tg = engine.threadgroupSize2D(pso, grid: (numHeads * headDim, seqLen)) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + + return output + } + + private func applyLConv1D( + input: MTLBuffer, + weights: AudioLayerWeights, + seqLen: Int, + cmdBuf: MTLCommandBuffer + ) throws -> MTLBuffer { + // Pre-layer norm + var current = try applyRMSNorm( + input: input, + weight: weights.lconv1dPreLayerNorm, + seqLen: seqLen, + hiddenSize: config.hiddenSize, + cmdBuf: cmdBuf + ) + + // Linear start: [seqLen, 1024] -> [seqLen, 2048] + let linearStart = try applyQuantizedLinear( + input: current, + weights: weights.lconv1dLinearStart, + seqLen: seqLen, + output: ffnBuffer, + cmdBuf: cmdBuf + ) + + // SiLU activation + let activated = try applySiLU(input: linearStart, count: seqLen * config.hiddenSize * 2, cmdBuf: cmdBuf) + + // Depthwise conv1d + let convOut = try applyDepthwiseConv1D( + input: activated, + weight: weights.lconv1dDepthwiseConv, + norm: weights.lconv1dConvNorm, + seqLen: seqLen, + channels: config.hiddenSize * 2, + kernelSize: config.convKernelSize, + cmdBuf: cmdBuf + ) + + // Linear end: [seqLen, 2048] -> [seqLen, 1024] + let output = try applyQuantizedLinear( + input: convOut, + weights: weights.lconv1dLinearEnd, + seqLen: seqLen, + output: tempBuffer, + cmdBuf: cmdBuf + ) + + return output + } + + private func applyDepthwiseConv1D( + input: MTLBuffer, + weight: MTLBuffer, + norm: MTLBuffer, + seqLen: Int, + channels: Int, + kernelSize: Int, + cmdBuf: MTLCommandBuffer + ) throws -> MTLBuffer { + // FIX: Use layerBuffer for audio layers + let output = layerBuffer + + let pso = try engine.pipeline(named: "audio_depthwise_conv1d") + let enc = cmdBuf.makeComputeCommandEncoder()! + enc.setComputePipelineState(pso) + + enc.setBuffer(input, offset: 0, index: 0) + enc.setBuffer(weight, offset: 0, index: 1) + enc.setBuffer(norm, offset: 0, index: 2) + enc.setBuffer(output, offset: 0, index: 3) + + var channels_ = UInt32(channels) + enc.setBytes(&channels_, length: 4, index: 4) + var kernelSize_ = UInt32(kernelSize) + enc.setBytes(&kernelSize_, length: 4, index: 5) + var seqLen_ = UInt32(seqLen) + enc.setBytes(&seqLen_, length: 4, index: 6) + + let grid = MTLSize(width: channels, height: seqLen, depth: 1) + let tg = engine.threadgroupSize2D(pso, grid: (channels, seqLen)) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + + return output + } + + private func applyFeedForward( + input: MTLBuffer, + weights: FeedForwardWeights, + seqLen: Int, + cmdBuf: MTLCommandBuffer + ) throws -> MTLBuffer { + // Pre-layer norm + var current = try applyRMSNorm( + input: input, + weight: weights.preLayerNorm, + seqLen: seqLen, + hiddenSize: config.hiddenSize, + cmdBuf: cmdBuf + ) + + // Layer 1: [seqLen, 1024] -> [seqLen, 4096] + let layer1 = try applyQuantizedLinear( + input: current, + weights: weights.ffwLayer1, + seqLen: seqLen, + output: ffnBuffer, + cmdBuf: cmdBuf + ) + + // SiLU activation + let activated = try applySiLU(input: layer1, count: seqLen * 4096, cmdBuf: cmdBuf) + + // Layer 2: [seqLen, 4096] -> [seqLen, 1024] + let output = try applyQuantizedLinear( + input: activated, + weights: weights.ffwLayer2, + seqLen: seqLen, + output: tempBuffer, + cmdBuf: cmdBuf + ) + + // Post-layer norm + return try applyRMSNorm( + input: output, + weight: weights.postLayerNorm, + seqLen: seqLen, + hiddenSize: config.hiddenSize, + cmdBuf: cmdBuf + ) + } + + private func applyRMSNorm( + input: MTLBuffer, + weight: MTLBuffer, + seqLen: Int, + hiddenSize: Int, + cmdBuf: MTLCommandBuffer + ) throws -> MTLBuffer { + // FIX: Use layerBuffer for audio layers to avoid tempBuffer conflict + let output = layerBuffer + + let pso = try engine.pipeline(named: "rms_norm_seq") + let enc = cmdBuf.makeComputeCommandEncoder()! + 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 N = UInt32(hiddenSize) + enc.setBytes(&N, length: 4, index: 3) + var eps = config.rmsNormEps + enc.setBytes(&eps, length: 4, index: 4) + var seqLen_ = UInt32(seqLen) + enc.setBytes(&seqLen_, length: 4, index: 5) + + let grid = MTLSize(width: hiddenSize, height: seqLen, depth: 1) + let tg = engine.threadgroupSize2D(pso, grid: (hiddenSize, seqLen)) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + + return output + } + + private func applyQuantizedLinear( + input: MTLBuffer, + weights: QuantizedWeights, + seqLen: Int, + output: MTLBuffer, + bias: MTLBuffer? = nil, + cmdBuf: MTLCommandBuffer + ) throws -> MTLBuffer { + let pso = try engine.pipeline(named: "quantized_matmul_seq") + let enc = cmdBuf.makeComputeCommandEncoder()! + enc.setComputePipelineState(pso) + + enc.setBuffer(input, offset: 0, index: 0) + enc.setBuffer(weights.weight, offset: 0, index: 1) + enc.setBuffer(weights.scales, offset: 0, index: 2) + enc.setBuffer(weights.biases, offset: 0, index: 3) + enc.setBuffer(bias, offset: 0, index: 4) + enc.setBuffer(output, offset: 0, index: 5) + + var inDim = UInt32(weights.inDim) + enc.setBytes(&inDim, length: 4, index: 6) + var outDim = UInt32(weights.outDim) + enc.setBytes(&outDim, length: 4, index: 7) + var hasBias = bias != nil + enc.setBytes(&hasBias, length: 1, index: 8) + var seqLen_ = UInt32(seqLen) + enc.setBytes(&seqLen_, length: 4, index: 9) + + let grid = MTLSize(width: weights.outDim, height: seqLen, depth: 1) + let tg = engine.threadgroupSize2D(pso, grid: (weights.outDim, seqLen)) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + + return output + } + + private func applySiLU(input: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer { + // FIX: Use layerBuffer for audio layers + let output = layerBuffer + + let pso = try engine.pipeline(named: "silu") + let enc = cmdBuf.makeComputeCommandEncoder()! + enc.setComputePipelineState(pso) + + enc.setBuffer(input, offset: 0, index: 0) + enc.setBuffer(output, offset: 0, index: 1) + + var count_ = UInt32(count) + enc.setBytes(&count_, length: 4, index: 2) + + let grid = MTLSize(width: count, height: 1, depth: 1) + let tg = engine.threadgroupSize1D(pso, count: count) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + + return output + } + + private func applyResidualAdd( + input: MTLBuffer, + add: MTLBuffer, + seqLen: Int, + hiddenSize: Int, + residualWeight: Float, + cmdBuf: MTLCommandBuffer + ) throws -> MTLBuffer { + // FIX: Use layerBuffer for audio layers + let output = layerBuffer + + let pso = try engine.pipeline(named: "residual_add") + let enc = cmdBuf.makeComputeCommandEncoder()! + enc.setComputePipelineState(pso) + + enc.setBuffer(input, offset: 0, index: 0) + enc.setBuffer(add, offset: 0, index: 1) + enc.setBuffer(output, offset: 0, index: 2) + + var count32 = UInt32(seqLen * hiddenSize) + enc.setBytes(&count32, length: 4, index: 3) + var weight = residualWeight + enc.setBytes(&weight, length: 4, index: 4) + + let count = seqLen * hiddenSize + let grid = MTLSize(width: count, height: 1, depth: 1) + let tg = engine.threadgroupSize1D(pso, count: count) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + + return output + } + + private func applyOutputProjection( + input: MTLBuffer, + seqLen: Int, + output: MTLBuffer, + cmdBuf: MTLCommandBuffer + ) throws { + _ = try applyQuantizedLinear( + input: input, + weights: weights.outputProj, + seqLen: seqLen, + output: output, + bias: weights.outputProjBias, + cmdBuf: cmdBuf + ) + } +} \ No newline at end of file diff --git a/Sources/MarkBase/Audio/AudioTower12B.swift b/Sources/MarkBase/Audio/AudioTower12B.swift new file mode 100644 index 0000000..cccbbb2 --- /dev/null +++ b/Sources/MarkBase/Audio/AudioTower12B.swift @@ -0,0 +1,129 @@ +import Metal + +public struct AudioConfig12B { + public let outputDim: Int + public let audioDim: Int + public let groupSize: Int + + public init(outputDim: Int = 3840, audioDim: Int = 640, groupSize: Int = 64) { + self.outputDim = outputDim + self.audioDim = audioDim + self.groupSize = groupSize + } +} + +public struct AudioWeights12B { + public let projectionWeight: MTLBuffer + public let projectionScales: MTLBuffer + public let projectionBiases: MTLBuffer + public let numGroups: Int + public let hasOutputBias: Bool + public let outputBias: MTLBuffer? + + public init(device: MTLDevice, + weightData: [UInt32], + scalesData: [Float], + biasesData: [Float], + numGroups: Int, + outputBias: [Float]? = nil) throws { + projectionWeight = device.makeBuffer(bytes: weightData, length: weightData.count * 4)! + projectionScales = device.makeBuffer(bytes: scalesData, length: scalesData.count * 4)! + projectionBiases = device.makeBuffer(bytes: biasesData, length: biasesData.count * 4)! + self.numGroups = numGroups + + if let bias = outputBias { + self.outputBias = device.makeBuffer(bytes: bias, length: bias.count * 4) + self.hasOutputBias = true + } else { + self.outputBias = nil + self.hasOutputBias = false + } + } +} + +public final class AudioTower12B { + public let config: AudioConfig12B + public let weights: AudioWeights12B + public let engine: MarkBaseEngine + + public let inDim: Int + public let outDim: Int + + public init(config: AudioConfig12B, engine: MarkBaseEngine, weights: AudioWeights12B) { + self.config = config + self.weights = weights + self.engine = engine + + self.inDim = config.audioDim + self.outDim = config.outputDim + } + + public func forward(inputBuffer: MTLBuffer, seqLen: Int, outputBuffer: MTLBuffer) throws { + let cmdBuf = engine.commandQueue.makeCommandBuffer()! + defer { cmdBuf.commit(); cmdBuf.waitUntilCompleted() } + + let pso = try engine.pipeline(named: "quantized_matmul_seq") + let enc = cmdBuf.makeComputeCommandEncoder()! + enc.setComputePipelineState(pso) + + enc.setBuffer(inputBuffer, offset: 0, index: 0) + enc.setBuffer(weights.projectionWeight, offset: 0, index: 1) + enc.setBuffer(weights.projectionScales, offset: 0, index: 2) + enc.setBuffer(weights.projectionBiases, offset: 0, index: 3) + + if let bias = weights.outputBias { + enc.setBuffer(bias, offset: 0, index: 4) + } else { + enc.setBuffer(weights.projectionBiases, offset: 0, index: 4) + } + + enc.setBuffer(outputBuffer, offset: 0, index: 5) + + var inD = UInt32(inDim) + enc.setBytes(&inD, length: 4, index: 6) + var outD = UInt32(outDim) + enc.setBytes(&outD, length: 4, index: 7) + var hasBias = weights.hasOutputBias + enc.setBytes(&hasBias, length: 1, index: 8) + var sl = UInt32(seqLen) + enc.setBytes(&sl, length: 4, index: 9) + + let grid = MTLSize(width: outDim, height: seqLen, depth: 1) + let tg = engine.threadgroupSize2D(pso, grid: (outDim, seqLen)) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + } + + public static func load(modelDir: String, engine: MarkBaseEngine) throws -> AudioTower12B { + let device = engine.device + let shardFile = "model-00002-of-00002.safetensors" + let reader = try SafeTensorsReader(path: "\(modelDir)/\(shardFile)") + + let weightData = try reader.readUint32(named: "embed_audio.embedding_projection.weight") + let scalesRaw = try reader.read(named: "embed_audio.embedding_projection.scales") + let scalesData = SafeTensorsReader.bf16ToFloat32(scalesRaw) + let biasesRaw = try reader.read(named: "embed_audio.embedding_projection.biases") + let biasesData = SafeTensorsReader.bf16ToFloat32(biasesRaw) + + let numWeights = weightData.count + let numScales = scalesData.count + + let audioDim = 640 + let packedInDim = audioDim / 8 + let outDim = numWeights / packedInDim + let numGroups = packedInDim / 8 + + let weights = try AudioWeights12B( + device: device, + weightData: weightData, + scalesData: scalesData, + biasesData: biasesData, + numGroups: numGroups, + outputBias: nil + ) + + let config = AudioConfig12B(outputDim: outDim, audioDim: audioDim, groupSize: 64) + print(" AudioTower12B: inDim=\(audioDim), outDim=\(outDim), numGroups=\(numGroups)") + return AudioTower12B(config: config, engine: engine, weights: weights) + } +} \ No newline at end of file diff --git a/Sources/MarkBase/Audio/AudioTowerE2B.swift b/Sources/MarkBase/Audio/AudioTowerE2B.swift new file mode 100644 index 0000000..4027e55 --- /dev/null +++ b/Sources/MarkBase/Audio/AudioTowerE2B.swift @@ -0,0 +1,609 @@ +import Metal + +// E2B audio tower uses bfloat16 weights (not quantized) +// Linear weights are full bfloat16, not uint32 packed + +public struct AudioLayerWeightsE2B { + public let normPreAttn: MTLBuffer + public let normPostAttn: MTLBuffer + public let normOut: MTLBuffer + + public let selfAttnQProjWeight: MTLBuffer + public let selfAttnKProjWeight: MTLBuffer + public let selfAttnVProjWeight: MTLBuffer + public let selfAttnPostWeight: MTLBuffer + public let selfAttnRelativeKProj: MTLBuffer + public let selfAttnPerDimScale: MTLBuffer + + public let lconv1dPreLayerNorm: MTLBuffer + public let lconv1dConvNorm: MTLBuffer + public let lconv1dDepthwiseConv: MTLBuffer + public let lconv1dLinearStartWeight: MTLBuffer + public let lconv1dLinearEndWeight: MTLBuffer + + public let feedForward1Layer1Weight: MTLBuffer + public let feedForward1Layer2Weight: MTLBuffer + public let feedForward1PreLayerNorm: MTLBuffer + public let feedForward1PostLayerNorm: MTLBuffer + + public let feedForward2Layer1Weight: MTLBuffer + public let feedForward2Layer2Weight: MTLBuffer + public let feedForward2PreLayerNorm: MTLBuffer + public let feedForward2PostLayerNorm: MTLBuffer + + private static func buffer(_ device: MTLDevice, _ floats: [String: [Float]], + _ key: String) throws -> MTLBuffer { + guard let f = floats[key] else { + throw WeightError.tensorNotFound(key) + } + guard let buf = device.makeBuffer(bytes: f, length: f.count * MemoryLayout.stride) else { + throw WeightError.tensorNotFound("Failed to create buffer for \(key)") + } + return buf + } + + public init(device: MTLDevice, layerIdx: Int, + floats: [String: [Float]]) throws { + let P = "audio_tower.layers.\(layerIdx)." + + normPreAttn = try Self.buffer(device, floats, P + "norm_pre_attn.weight") + normPostAttn = try Self.buffer(device, floats, P + "norm_post_attn.weight") + normOut = try Self.buffer(device, floats, P + "norm_out.weight") + + // Attention projections - use linear.weight suffix for E2B + selfAttnQProjWeight = try Self.buffer(device, floats, P + "self_attn.q_proj.linear.weight") + selfAttnKProjWeight = try Self.buffer(device, floats, P + "self_attn.k_proj.linear.weight") + selfAttnVProjWeight = try Self.buffer(device, floats, P + "self_attn.v_proj.linear.weight") + selfAttnPostWeight = try Self.buffer(device, floats, P + "self_attn.post.linear.weight") + selfAttnRelativeKProj = try Self.buffer(device, floats, P + "self_attn.relative_k_proj.weight") + selfAttnPerDimScale = try Self.buffer(device, floats, P + "self_attn.per_dim_scale") + + // LConv1D + lconv1dPreLayerNorm = try Self.buffer(device, floats, P + "lconv1d.pre_layer_norm.weight") + lconv1dConvNorm = try Self.buffer(device, floats, P + "lconv1d.conv_norm.weight") + lconv1dDepthwiseConv = try Self.buffer(device, floats, P + "lconv1d.depthwise_conv1d.weight") + lconv1dLinearStartWeight = try Self.buffer(device, floats, P + "lconv1d.linear_start.linear.weight") + lconv1dLinearEndWeight = try Self.buffer(device, floats, P + "lconv1d.linear_end.linear.weight") + + // FeedForward 1 + feedForward1Layer1Weight = try Self.buffer(device, floats, P + "feed_forward1.ffw_layer_1.linear.weight") + feedForward1Layer2Weight = try Self.buffer(device, floats, P + "feed_forward1.ffw_layer_2.linear.weight") + feedForward1PreLayerNorm = try Self.buffer(device, floats, P + "feed_forward1.pre_layer_norm.weight") + feedForward1PostLayerNorm = try Self.buffer(device, floats, P + "feed_forward1.post_layer_norm.weight") + + // FeedForward 2 + feedForward2Layer1Weight = try Self.buffer(device, floats, P + "feed_forward2.ffw_layer_1.linear.weight") + feedForward2Layer2Weight = try Self.buffer(device, floats, P + "feed_forward2.ffw_layer_2.linear.weight") + feedForward2PreLayerNorm = try Self.buffer(device, floats, P + "feed_forward2.pre_layer_norm.weight") + feedForward2PostLayerNorm = try Self.buffer(device, floats, P + "feed_forward2.post_layer_norm.weight") + } +} + +public struct AudioWeightsE2B { + public let subsampleConvLayer0: SubsampleConvLayer + public let subsampleConvLayer1: SubsampleConvLayer + public let inputProjLinearWeight: MTLBuffer + + public let outputProjWeight: MTLBuffer + public let outputProjBias: MTLBuffer + + public let layers: [AudioLayerWeightsE2B] + + public init(device: MTLDevice, config: AudioConfig, + floats: [String: [Float]]) throws { + let P = "audio_tower." + + subsampleConvLayer0 = SubsampleConvLayer( + convWeight: try Self.buffer(device, floats, P + "subsample_conv_projection.layer0.conv.weight"), + normWeight: try Self.buffer(device, floats, P + "subsample_conv_projection.layer0.norm.weight") + ) + subsampleConvLayer1 = SubsampleConvLayer( + convWeight: try Self.buffer(device, floats, P + "subsample_conv_projection.layer1.conv.weight"), + normWeight: try Self.buffer(device, floats, P + "subsample_conv_projection.layer1.norm.weight") + ) + inputProjLinearWeight = try Self.buffer(device, floats, P + "subsample_conv_projection.input_proj_linear.weight") + + outputProjWeight = try Self.buffer(device, floats, P + "output_proj.weight") + outputProjBias = try Self.buffer(device, floats, P + "output_proj.bias") + + var loadedLayers: [AudioLayerWeightsE2B] = [] + for i in 0.. MTLBuffer { + guard let f = floats[key] else { + throw WeightError.tensorNotFound(key) + } + guard let buf = device.makeBuffer(bytes: f, length: f.count * MemoryLayout.stride) else { + throw WeightError.tensorNotFound("Failed to create buffer for \(key)") + } + return buf + } +} + +// E2B AudioTower - uses float32 weights (bfloat16 converted to float32) +public final class AudioTowerE2B { + public let config: AudioConfig + public let engine: MarkBaseEngine + public let weights: AudioWeightsE2B + + private var normBuffer: MTLBuffer + private var qBuffer: MTLBuffer + private var kBuffer: MTLBuffer + private var vBuffer: MTLBuffer + private var attnOutBuffer: MTLBuffer + private var ffnBuffer: MTLBuffer + private var tempBuffer: MTLBuffer + private var subsampleBuf: MTLBuffer + + public init(config: AudioConfig, engine: MarkBaseEngine, weights: AudioWeightsE2B) throws { + self.config = config + self.engine = engine + self.weights = weights + + let device = engine.device + let maxSeqLen = 4096 + let hiddenSize = config.hiddenSize + + normBuffer = device.makeBuffer(length: hiddenSize * maxSeqLen * 4)! + qBuffer = device.makeBuffer(length: hiddenSize * maxSeqLen * 4)! + kBuffer = device.makeBuffer(length: hiddenSize * maxSeqLen * 4)! + vBuffer = device.makeBuffer(length: hiddenSize * maxSeqLen * 4)! + attnOutBuffer = device.makeBuffer(length: hiddenSize * maxSeqLen * 4)! + ffnBuffer = device.makeBuffer(length: 4096 * maxSeqLen * 4)! + tempBuffer = device.makeBuffer(length: max(hiddenSize, 4096) * maxSeqLen * 4)! + subsampleBuf = device.makeBuffer(length: max(hiddenSize, 128 * 64) * maxSeqLen * 4)! + } + + public func forward(inputBuffer: MTLBuffer, seqLen: Int, outputBuffer: MTLBuffer) throws { + var current = inputBuffer + var currentLen = seqLen + + let cmdBuf = engine.commandQueue.makeCommandBuffer()! + + // 1. Subsample conv + let (projInput, projLen) = try applySubsampleConv( + melInput: current, nMels: 128, seqLen: currentLen, cmdBuf: cmdBuf + ) + + // 2. Input projection + current = try applyFloatLinear(input: projInput, weight: weights.inputProjLinearWeight, + seqLen: projLen, inDim: 1024, outDim: 1024, cmdBuf: cmdBuf) + currentLen = projLen + + // 3. Audio layers + for layerWeights in weights.layers { + current = try applyLayer(input: current, weights: layerWeights, + seqLen: currentLen, cmdBuf: cmdBuf) + } + + // 4. Output projection + try applyOutputProjection(input: current, seqLen: currentLen, output: outputBuffer, cmdBuf: cmdBuf) + + cmdBuf.commit() + cmdBuf.waitUntilCompleted() + } + + private func applySubsampleConv(melInput: MTLBuffer, nMels: Int, seqLen: Int, + cmdBuf: MTLCommandBuffer) throws -> (MTLBuffer, Int) { + let chwInput = try transposeMelToCHW(input: melInput, nMels: nMels, seqLen: seqLen, cmdBuf: cmdBuf) + + let layer0Out = try applyConv2DLayer(input: chwInput, inCh: 1, height: nMels, width: seqLen, + convWeight: weights.subsampleConvLayer0.convWeight, + normWeight: weights.subsampleConvLayer0.normWeight, + outChannels: 128, outputBuffer: tempBuffer, cmdBuf: cmdBuf) + let h1 = (nMels + 1) / 2 + let w1 = (seqLen + 1) / 2 + + let layer1Out = try applyConv2DLayer(input: layer0Out, inCh: 128, height: h1, width: w1, + convWeight: weights.subsampleConvLayer1.convWeight, + normWeight: weights.subsampleConvLayer1.normWeight, + outChannels: 32, outputBuffer: subsampleBuf, cmdBuf: cmdBuf) + let h2 = (h1 + 1) / 2 + let w2 = (w1 + 1) / 2 + + let flatOutput = try flattenCHW(input: layer1Out, C: 32, H: h2, W: w2, + outputBuffer: tempBuffer, cmdBuf: cmdBuf) + return (flatOutput, w2) + } + + private func transposeMelToCHW(input: MTLBuffer, nMels: Int, seqLen: Int, + cmdBuf: MTLCommandBuffer) throws -> MTLBuffer { + let output = subsampleBuf + let pso = try engine.pipeline(named: "transpose_2d") + let enc = cmdBuf.makeComputeCommandEncoder()! + enc.setComputePipelineState(pso) + enc.setBuffer(input, offset: 0, index: 0) + enc.setBuffer(output, offset: 0, index: 1) + var rows = UInt32(nMels) + enc.setBytes(&rows, length: 4, index: 2) + var cols = UInt32(seqLen) + enc.setBytes(&cols, length: 4, index: 3) + let grid = MTLSize(width: seqLen, height: nMels, depth: 1) + let tg = engine.threadgroupSize2D(pso, grid: (seqLen, nMels)) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + return output + } + + private func applyConv2DLayer(input: MTLBuffer, inCh: Int, height: Int, width: Int, + convWeight: MTLBuffer, normWeight: MTLBuffer, + outChannels: Int, outputBuffer: MTLBuffer, + cmdBuf: MTLCommandBuffer) throws -> MTLBuffer { + let pso = try engine.pipeline(named: "audio_subsample_conv_2d") + let enc = cmdBuf.makeComputeCommandEncoder()! + enc.setComputePipelineState(pso) + enc.setBuffer(input, offset: 0, index: 0) + enc.setBuffer(convWeight, offset: 0, index: 1) + enc.setBuffer(normWeight, offset: 0, index: 2) + enc.setBuffer(outputBuffer, offset: 0, index: 3) + var inCh_ = UInt32(inCh) + enc.setBytes(&inCh_, length: 4, index: 4) + var outCh_ = UInt32(outChannels) + enc.setBytes(&outCh_, length: 4, index: 5) + var h_ = UInt32(height) + enc.setBytes(&h_, length: 4, index: 6) + var w_ = UInt32(width) + enc.setBytes(&w_, length: 4, index: 7) + let outH = (height + 1) / 2 + let outW = (width + 1) / 2 + let grid = MTLSize(width: outChannels, height: outH, depth: outW) + let tg = MTLSize(width: 8, height: 8, depth: 4) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + return outputBuffer + } + + private func flattenCHW(input: MTLBuffer, C: Int, H: Int, W: Int, + outputBuffer: MTLBuffer, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer { + let pso = try engine.pipeline(named: "audio_flatten_chw") + let enc = cmdBuf.makeComputeCommandEncoder()! + enc.setComputePipelineState(pso) + enc.setBuffer(input, offset: 0, index: 0) + enc.setBuffer(outputBuffer, offset: 0, index: 1) + var C_ = UInt32(C) + enc.setBytes(&C_, length: 4, index: 2) + var H_ = UInt32(H) + enc.setBytes(&H_, length: 4, index: 3) + var W_ = UInt32(W) + enc.setBytes(&W_, length: 4, index: 4) + let grid = MTLSize(width: C * H, height: W, depth: 1) + let tg = engine.threadgroupSize2D(pso, grid: (C * H, W)) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + return outputBuffer + } + + private func applyFloatLinear(input: MTLBuffer, weight: MTLBuffer, seqLen: Int, + inDim: Int, outDim: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer { + let output = tempBuffer + let pso = try engine.pipeline(named: "audio_linear_seq") + let enc = cmdBuf.makeComputeCommandEncoder()! + enc.setComputePipelineState(pso) + enc.setBuffer(input, offset: 0, index: 0) + enc.setBuffer(weight, offset: 0, index: 1) + enc.setBuffer(nil, offset: 0, index: 2) // No bias + enc.setBuffer(output, offset: 0, index: 3) + var inF = UInt32(inDim) + enc.setBytes(&inF, length: 4, index: 4) + var outF = UInt32(outDim) + enc.setBytes(&outF, length: 4, index: 5) + var hasBias = false + enc.setBytes(&hasBias, length: 1, index: 6) + var seqLen_ = UInt32(seqLen) + enc.setBytes(&seqLen_, length: 4, index: 7) + let grid = MTLSize(width: outDim, height: seqLen, depth: 1) + let tg = engine.threadgroupSize2D(pso, grid: (outDim, seqLen)) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + return output + } + + private func applyLayer(input: MTLBuffer, weights: AudioLayerWeightsE2B, + seqLen: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer { + var current = input + + // 1. Norm pre-attn + current = try applyRMSNorm(input: current, weight: weights.normPreAttn, + seqLen: seqLen, hiddenSize: config.hiddenSize, cmdBuf: cmdBuf) + + // 2. Self-attention + let q = try applyFloatLinear(input: current, weight: weights.selfAttnQProjWeight, + seqLen: seqLen, inDim: config.hiddenSize, outDim: config.hiddenSize, cmdBuf: cmdBuf) + let k = try applyFloatLinear(input: current, weight: weights.selfAttnKProjWeight, + seqLen: seqLen, inDim: config.hiddenSize, outDim: config.hiddenSize, cmdBuf: cmdBuf) + let v = try applyFloatLinear(input: current, weight: weights.selfAttnVProjWeight, + seqLen: seqLen, inDim: config.hiddenSize, outDim: config.hiddenSize, cmdBuf: cmdBuf) + + let attnOut = try applyAudioAttention(q: q, k: k, v: v, + relativeKProj: weights.selfAttnRelativeKProj, + perDimScale: weights.selfAttnPerDimScale, + seqLen: seqLen, cmdBuf: cmdBuf) + + let post = try applyFloatLinear(input: attnOut, weight: weights.selfAttnPostWeight, + seqLen: seqLen, inDim: config.hiddenSize, outDim: config.hiddenSize, cmdBuf: cmdBuf) + + // 3. Residual + norm + current = try applyResidualAdd(input: input, add: post, seqLen: seqLen, + hiddenSize: config.hiddenSize, cmdBuf: cmdBuf) + current = try applyRMSNorm(input: current, weight: weights.normPostAttn, + seqLen: seqLen, hiddenSize: config.hiddenSize, cmdBuf: cmdBuf) + + // 4. LConv1D + let lconvOut = try applyLConv1D(input: current, weights: weights, seqLen: seqLen, cmdBuf: cmdBuf) + current = try applyResidualAdd(input: current, add: lconvOut, seqLen: seqLen, + hiddenSize: config.hiddenSize, cmdBuf: cmdBuf) + + // 5. FeedForward 1 + let ff1Out = try applyFeedForward(input: current, + layer1Weight: weights.feedForward1Layer1Weight, + layer2Weight: weights.feedForward1Layer2Weight, + preNorm: weights.feedForward1PreLayerNorm, + postNorm: weights.feedForward1PostLayerNorm, + seqLen: seqLen, cmdBuf: cmdBuf) + current = try applyResidualAdd(input: current, add: ff1Out, seqLen: seqLen, + hiddenSize: config.hiddenSize, cmdBuf: cmdBuf) + + // 6. FeedForward 2 + let ff2Out = try applyFeedForward(input: current, + layer1Weight: weights.feedForward2Layer1Weight, + layer2Weight: weights.feedForward2Layer2Weight, + preNorm: weights.feedForward2PreLayerNorm, + postNorm: weights.feedForward2PostLayerNorm, + seqLen: seqLen, cmdBuf: cmdBuf) + current = try applyResidualAdd(input: current, add: ff2Out, seqLen: seqLen, + hiddenSize: config.hiddenSize, cmdBuf: cmdBuf) + current = try applyRMSNorm(input: current, weight: weights.normOut, + seqLen: seqLen, hiddenSize: config.hiddenSize, cmdBuf: cmdBuf) + + return current + } + + private func applyRMSNorm(input: MTLBuffer, weight: MTLBuffer, seqLen: Int, + hiddenSize: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer { + let output = tempBuffer + let pso = try engine.pipeline(named: "rms_norm_seq") + let enc = cmdBuf.makeComputeCommandEncoder()! + 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 N = UInt32(hiddenSize) + enc.setBytes(&N, length: 4, index: 3) + var eps = config.rmsNormEps + enc.setBytes(&eps, length: 4, index: 4) + var seqLen_ = UInt32(seqLen) + enc.setBytes(&seqLen_, length: 4, index: 5) + let grid = MTLSize(width: hiddenSize, height: seqLen, depth: 1) + let tg = engine.threadgroupSize2D(pso, grid: (hiddenSize, seqLen)) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + return output + } + + private func applyAudioAttention(q: MTLBuffer, k: MTLBuffer, v: MTLBuffer, + relativeKProj: MTLBuffer, perDimScale: MTLBuffer, + seqLen: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer { + let output = attnOutBuffer + let pso = try engine.pipeline(named: "audio_attention_full") + let enc = cmdBuf.makeComputeCommandEncoder()! + 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(relativeKProj, offset: 0, index: 3) + enc.setBuffer(perDimScale, offset: 0, index: 4) + enc.setBuffer(output, offset: 0, index: 5) + var seqLen_ = UInt32(seqLen) + enc.setBytes(&seqLen_, length: 4, index: 6) + var numHeads = UInt32(config.numAttentionHeads) + enc.setBytes(&numHeads, length: 4, index: 7) + var headDim = UInt32(config.headDim) + enc.setBytes(&headDim, length: 4, index: 8) + var contextLeft = UInt32(config.attentionContextLeft) + enc.setBytes(&contextLeft, length: 4, index: 9) + var logitCap = config.attentionLogitCap + enc.setBytes(&logitCap, length: 4, index: 10) + let grid = MTLSize(width: config.numAttentionHeads * config.headDim, height: seqLen, depth: 1) + let tg = engine.threadgroupSize2D(pso, grid: (config.numAttentionHeads * config.headDim, seqLen)) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + return output + } + + private func applyLConv1D(input: MTLBuffer, weights: AudioLayerWeightsE2B, + seqLen: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer { + var current = try applyRMSNorm(input: input, weight: weights.lconv1dPreLayerNorm, + seqLen: seqLen, hiddenSize: config.hiddenSize, cmdBuf: cmdBuf) + + let linearStart = try applyFloatLinear(input: current, weight: weights.lconv1dLinearStartWeight, + seqLen: seqLen, inDim: config.hiddenSize, + outDim: config.hiddenSize * 2, cmdBuf: cmdBuf) + let activated = try applySiLU(input: linearStart, count: seqLen * config.hiddenSize * 2, cmdBuf: cmdBuf) + let convOut = try applyDepthwiseConv1D(input: activated, weight: weights.lconv1dDepthwiseConv, + norm: weights.lconv1dConvNorm, seqLen: seqLen, + channels: config.hiddenSize * 2, kernelSize: config.convKernelSize, + cmdBuf: cmdBuf) + let linearEnd = try applyFloatLinear(input: convOut, weight: weights.lconv1dLinearEndWeight, + seqLen: seqLen, inDim: config.hiddenSize * 2, + outDim: config.hiddenSize, cmdBuf: cmdBuf) + return linearEnd + } + + private func applyDepthwiseConv1D(input: MTLBuffer, weight: MTLBuffer, norm: MTLBuffer, + seqLen: Int, channels: Int, kernelSize: Int, + cmdBuf: MTLCommandBuffer) throws -> MTLBuffer { + let output = tempBuffer + let pso = try engine.pipeline(named: "audio_depthwise_conv1d") + let enc = cmdBuf.makeComputeCommandEncoder()! + enc.setComputePipelineState(pso) + enc.setBuffer(input, offset: 0, index: 0) + enc.setBuffer(weight, offset: 0, index: 1) + enc.setBuffer(norm, offset: 0, index: 2) + enc.setBuffer(output, offset: 0, index: 3) + var channels_ = UInt32(channels) + enc.setBytes(&channels_, length: 4, index: 4) + var kernelSize_ = UInt32(kernelSize) + enc.setBytes(&kernelSize_, length: 4, index: 5) + var seqLen_ = UInt32(seqLen) + enc.setBytes(&seqLen_, length: 4, index: 6) + let grid = MTLSize(width: channels, height: seqLen, depth: 1) + let tg = engine.threadgroupSize2D(pso, grid: (channels, seqLen)) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + return output + } + + private func applyFeedForward(input: MTLBuffer, layer1Weight: MTLBuffer, layer2Weight: MTLBuffer, + preNorm: MTLBuffer, postNorm: MTLBuffer, + seqLen: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer { + var current = try applyRMSNorm(input: input, weight: preNorm, + seqLen: seqLen, hiddenSize: config.hiddenSize, cmdBuf: cmdBuf) + let layer1 = try applyFloatLinear(input: current, weight: layer1Weight, + seqLen: seqLen, inDim: config.hiddenSize, outDim: 4096, cmdBuf: cmdBuf) + let activated = try applySiLU(input: layer1, count: seqLen * 4096, cmdBuf: cmdBuf) + let layer2 = try applyFloatLinear(input: activated, weight: layer2Weight, + seqLen: seqLen, inDim: 4096, outDim: config.hiddenSize, cmdBuf: cmdBuf) + return try applyRMSNorm(input: layer2, weight: postNorm, + seqLen: seqLen, hiddenSize: config.hiddenSize, cmdBuf: cmdBuf) + } + + private func applySiLU(input: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer { + let output = tempBuffer + let pso = try engine.pipeline(named: "silu") + let enc = cmdBuf.makeComputeCommandEncoder()! + enc.setComputePipelineState(pso) + enc.setBuffer(input, offset: 0, index: 0) + enc.setBuffer(output, offset: 0, index: 1) + var count_ = UInt32(count) + enc.setBytes(&count_, length: 4, index: 2) + let grid = MTLSize(width: count, height: 1, depth: 1) + let tg = engine.threadgroupSize1D(pso, count: count) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + return output + } + + private func applyResidualAdd(input: MTLBuffer, add: MTLBuffer, seqLen: Int, + hiddenSize: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer { + let output = tempBuffer + let pso = try engine.pipeline(named: "residual_add") + let enc = cmdBuf.makeComputeCommandEncoder()! + enc.setComputePipelineState(pso) + enc.setBuffer(input, offset: 0, index: 0) + enc.setBuffer(add, offset: 0, index: 1) + enc.setBuffer(output, offset: 0, index: 2) + var count32 = UInt32(seqLen * hiddenSize) + enc.setBytes(&count32, length: 4, index: 3) + var weight = config.residualWeight + enc.setBytes(&weight, length: 4, index: 4) + let count = seqLen * hiddenSize + let grid = MTLSize(width: count, height: 1, depth: 1) + let tg = engine.threadgroupSize1D(pso, count: count) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + return output + } + + private func applyOutputProjection(input: MTLBuffer, seqLen: Int, output: MTLBuffer, + cmdBuf: MTLCommandBuffer) throws { + let pso = try engine.pipeline(named: "audio_linear_seq") + let enc = cmdBuf.makeComputeCommandEncoder()! + enc.setComputePipelineState(pso) + enc.setBuffer(input, offset: 0, index: 0) + enc.setBuffer(weights.outputProjWeight, offset: 0, index: 1) + enc.setBuffer(weights.outputProjBias, offset: 0, index: 2) + enc.setBuffer(output, offset: 0, index: 3) + var inF = UInt32(config.hiddenSize) + enc.setBytes(&inF, length: 4, index: 4) + var outF = UInt32(config.outputProjDims) + enc.setBytes(&outF, length: 4, index: 5) + var hasBias = true + enc.setBytes(&hasBias, length: 1, index: 6) + var seqLen_ = UInt32(seqLen) + enc.setBytes(&seqLen_, length: 4, index: 7) + let grid = MTLSize(width: config.outputProjDims, height: seqLen, depth: 1) + let tg = engine.threadgroupSize2D(pso, grid: (config.outputProjDims, seqLen)) + enc.dispatchThreads(grid, threadsPerThreadgroup: tg) + enc.endEncoding() + } +} + +// E2B audio tower loader function +func loadAudioTowerE2B(reader: SafeTensorsReader, config: AudioConfig, + engine: MarkBaseEngine) throws -> AudioTowerE2B { + print("Loading E2B Audio Tower with preload optimization...") + let startTime = Date() + + // Collect all audio tensor descriptors + let audioPrefix = "audio_tower." + let audioDescriptors = reader.allDescriptors().filter { + $0.name.hasPrefix(audioPrefix) + } + + print(" Found \(audioDescriptors.count) audio tensors") + + // Parallel preload all audio tensors + let dispatchGroup = DispatchGroup() + let loadQueue = DispatchQueue(label: "audio-preload-e2b", attributes: .concurrent) + var loadedData: [Data?] = Array(repeating: nil, count: audioDescriptors.count) + var loadErrors: [Error?] = Array(repeating: nil, count: audioDescriptors.count) + + for (idx, desc) in audioDescriptors.enumerated() { + dispatchGroup.enter() + loadQueue.async { + do { + let data = try reader.read(tensor: desc) + loadedData[idx] = data + } catch { + loadErrors[idx] = error + } + dispatchGroup.leave() + } + } + + dispatchGroup.wait() + + // Check for errors + for (idx, error) in loadErrors.enumerated() { + if let err = error { + throw WeightError.readFailed("Failed to preload audio tensor \(audioDescriptors[idx].name): \(err)") + } + } + + let preloadTime = Date().timeIntervalSince(startTime) * 1000 + print(" ✓ Parallel preloaded \(audioDescriptors.count) audio tensors in \(String(format: "%.1f", preloadTime))ms") + + // Convert to floats dictionary + var floats: [String: [Float]] = [:] + + for (idx, desc) in audioDescriptors.enumerated() { + guard let data = loadedData[idx] else { continue } + let name = desc.name + switch desc.dtype { + case .bf16: + floats[name] = SafeTensorsReader.bf16ToFloat32(data) + case .f32: + floats[name] = data.withUnsafeBytes { + Array($0.assumingMemoryBound(to: Float.self)) + } + default: + break + } + } + + guard !floats.isEmpty else { + throw WeightError.tensorNotFound("Audio tower tensors") + } + + let weights = try AudioWeightsE2B(device: engine.device, config: config, floats: floats) + + let totalTime = Date().timeIntervalSince(startTime) * 1000 + print(" ✓ E2B Audio Tower loaded in \(String(format: "%.1f", totalTime))ms") + + return try AudioTowerE2B(config: config, engine: engine, weights: weights) +} diff --git a/Sources/MarkBase/Audio/AudioWeights.swift b/Sources/MarkBase/Audio/AudioWeights.swift new file mode 100644 index 0000000..39a2e5b --- /dev/null +++ b/Sources/MarkBase/Audio/AudioWeights.swift @@ -0,0 +1,209 @@ +import Metal +import Foundation + +public final class AudioWeights { + public let subsampleConvLayer0: SubsampleConvLayer + public let subsampleConvLayer1: SubsampleConvLayer + public let inputProjLinearWeight: MTLBuffer // Float32, not quantized + + public let outputProj: QuantizedWeights + public let outputProjBias: MTLBuffer + + public let layers: [AudioLayerWeights] + + public init(device: MTLDevice, config: AudioConfig, + tensors: [String: Data], floats: [String: [Float]], + descriptors: [String: TensorDescriptor]) throws { + let P = "audio_tower." + + subsampleConvLayer0 = SubsampleConvLayer( + convWeight: try Self.buffer(device, floats, P + "subsample_conv_projection.layer0.conv.weight"), + normWeight: try Self.buffer(device, floats, P + "subsample_conv_projection.layer0.norm.weight") + ) + + subsampleConvLayer1 = SubsampleConvLayer( + convWeight: try Self.buffer(device, floats, P + "subsample_conv_projection.layer1.conv.weight"), + normWeight: try Self.buffer(device, floats, P + "subsample_conv_projection.layer1.norm.weight") + ) + + inputProjLinearWeight = try Self.buffer(device, floats, P + "subsample_conv_projection.input_proj_linear.weight") + + outputProj = try Self.loadQuantized(device: device, tensors: tensors, floats: floats, + descriptors: descriptors, + name: P + "output_proj") + outputProjBias = try Self.buffer(device, floats, P + "output_proj.bias") + + var loadedLayers: [AudioLayerWeights] = [] + for i in 0.. MTLBuffer { + guard let f = floats[key] else { + throw WeightError.tensorNotFound(key) + } + guard let buf = device.makeBuffer(bytes: f, length: f.count * MemoryLayout.stride) else { + throw WeightError.tensorNotFound("Failed to create buffer for \(key)") + } + return buf + } + + static func loadQuantized(device: MTLDevice, tensors: [String: Data], + floats: [String: [Float]], + descriptors: [String: TensorDescriptor], + name: String) throws -> QuantizedWeights { + let wName = name + ".weight" + let sName = name + ".scales" + let bName = name + ".biases" + + guard let wData = tensors[wName], + let sFloats = floats[sName], + let bFloats = floats[bName], + let wDesc = descriptors[wName], + let sDesc = descriptors[sName] else { + throw WeightError.tensorNotFound(name) + } + + // Dimensions from descriptors: + // weight: [outDim, inDim/8] (U32 packed, 8 values per U32) + // scales: [outDim, numGroups] where numGroups = inDim / groupSize + let outDim = wDesc.shape[0] + let numGroups = sDesc.shape[1] + let groupSize = 64 // Audio uses fixed group_size=64 + let inDim = numGroups * groupSize + + guard let wBuf = device.makeBuffer(bytes: (wData as NSData).bytes, length: wData.count, + options: .storageModeShared) else { + throw WeightError.bufferCreationFailed(wName) + } + guard let sBuf = device.makeBuffer(bytes: sFloats, length: sFloats.count * MemoryLayout.stride, + options: .storageModeShared) else { + throw WeightError.bufferCreationFailed(sName) + } + guard let bBuf = device.makeBuffer(bytes: bFloats, length: bFloats.count * MemoryLayout.stride, + options: .storageModeShared) else { + throw WeightError.bufferCreationFailed(bName) + } + + return QuantizedWeights(weight: wBuf, scales: sBuf, biases: bBuf, + inDim: inDim, outDim: outDim, bits: 4, groupSize: groupSize) + } +} + +public struct SubsampleConvLayer { + public let convWeight: MTLBuffer + public let normWeight: MTLBuffer +} + +public struct AudioLayerWeights { + public let normPreAttn: MTLBuffer + public let normPostAttn: MTLBuffer + public let normOut: MTLBuffer + + public let selfAttnQProj: QuantizedWeights + public let selfAttnKProj: QuantizedWeights + public let selfAttnVProj: QuantizedWeights + public let selfAttnPost: QuantizedWeights + public let selfAttnRelativeKProj: MTLBuffer + public let selfAttnPerDimScale: MTLBuffer + + public let lconv1dPreLayerNorm: MTLBuffer + public let lconv1dConvNorm: MTLBuffer + public let lconv1dDepthwiseConv: MTLBuffer + public let lconv1dLinearStart: QuantizedWeights + public let lconv1dLinearEnd: QuantizedWeights + + public let feedForward1: FeedForwardWeights + public let feedForward2: FeedForwardWeights + + private static func buffer(_ device: MTLDevice, _ floats: [String: [Float]], + _ key: String) throws -> MTLBuffer { + guard let f = floats[key] else { + throw WeightError.tensorNotFound(key) + } + guard let buf = device.makeBuffer(bytes: f, length: f.count * MemoryLayout.stride) else { + throw WeightError.tensorNotFound("Failed to create buffer for \(key)") + } + return buf + } + + public init(device: MTLDevice, layerIdx: Int, + tensors: [String: Data], floats: [String: [Float]], + descriptors: [String: TensorDescriptor]) throws { + let P = "audio_tower.layers.\(layerIdx)." + + normPreAttn = try Self.buffer(device, floats, P + "norm_pre_attn.weight") + normPostAttn = try Self.buffer(device, floats, P + "norm_post_attn.weight") + normOut = try Self.buffer(device, floats, P + "norm_out.weight") + + selfAttnQProj = try AudioWeights.loadQuantized(device: device, tensors: tensors, floats: floats, + descriptors: descriptors, + name: P + "self_attn.q_proj") + selfAttnKProj = try AudioWeights.loadQuantized(device: device, tensors: tensors, floats: floats, + descriptors: descriptors, + name: P + "self_attn.k_proj") + selfAttnVProj = try AudioWeights.loadQuantized(device: device, tensors: tensors, floats: floats, + descriptors: descriptors, + name: P + "self_attn.v_proj") + selfAttnPost = try AudioWeights.loadQuantized(device: device, tensors: tensors, floats: floats, + descriptors: descriptors, + name: P + "self_attn.post") + + selfAttnRelativeKProj = try Self.buffer(device, floats, P + "self_attn.relative_k_proj.weight") + selfAttnPerDimScale = try Self.buffer(device, floats, P + "self_attn.per_dim_scale") + + lconv1dPreLayerNorm = try Self.buffer(device, floats, P + "lconv1d.pre_layer_norm.weight") + lconv1dConvNorm = try Self.buffer(device, floats, P + "lconv1d.conv_norm.weight") + lconv1dDepthwiseConv = try Self.buffer(device, floats, P + "lconv1d.depthwise_conv1d.weight") + + lconv1dLinearStart = try AudioWeights.loadQuantized(device: device, tensors: tensors, floats: floats, + descriptors: descriptors, + name: P + "lconv1d.linear_start") + lconv1dLinearEnd = try AudioWeights.loadQuantized(device: device, tensors: tensors, floats: floats, + descriptors: descriptors, + name: P + "lconv1d.linear_end") + + feedForward1 = try FeedForwardWeights(device: device, prefix: P + "feed_forward1", + tensors: tensors, floats: floats, + descriptors: descriptors) + feedForward2 = try FeedForwardWeights(device: device, prefix: P + "feed_forward2", + tensors: tensors, floats: floats, + descriptors: descriptors) + } +} + +public struct FeedForwardWeights { + public let preLayerNorm: MTLBuffer + public let postLayerNorm: MTLBuffer + public let ffwLayer1: QuantizedWeights + public let ffwLayer2: QuantizedWeights + + public init(device: MTLDevice, prefix: String, + tensors: [String: Data], floats: [String: [Float]], + descriptors: [String: TensorDescriptor]) throws { + let b = { (key: String) throws -> MTLBuffer in + guard let f = floats[key] else { throw WeightError.tensorNotFound(key) } + guard let buf = device.makeBuffer(bytes: f, length: f.count * MemoryLayout.stride) else { + throw WeightError.tensorNotFound("Failed to create buffer for \(key)") + } + return buf + } + + preLayerNorm = try b(prefix + ".pre_layer_norm.weight") + postLayerNorm = try b(prefix + ".post_layer_norm.weight") + + ffwLayer1 = try AudioWeights.loadQuantized(device: device, tensors: tensors, floats: floats, + descriptors: descriptors, + name: prefix + ".ffw_layer_1") + ffwLayer2 = try AudioWeights.loadQuantized(device: device, tensors: tensors, floats: floats, + descriptors: descriptors, + name: prefix + ".ffw_layer_2") + } +} diff --git a/Sources/MarkBase/BatchGeneration.swift b/Sources/MarkBase/BatchGeneration.swift new file mode 100644 index 0000000..c960eb1 --- /dev/null +++ b/Sources/MarkBase/BatchGeneration.swift @@ -0,0 +1,244 @@ +import Metal + +// Batch generation extension for E4BModel +// Goal: Generate multiple tokens in one pass (reduce kernel dispatches) + +extension E4BModel { + + /// Batch forward pass - process multiple tokens at once + /// Reduces kernel dispatches from 854*N → 854 (for N tokens) + /// Expected improvement: ~8x for batch generation + public func forwardBatch(tokenIds: [Int], positions: [Int]) throws -> [[Float]] { + guard tokenIds.count == positions.count else { + throw NSError(domain: "Batch", code: -1, + userInfo: [NSLocalizedDescriptionKey: "tokenIds and positions must have same count"]) + } + + let batchSize = tokenIds.count + if batchSize == 0 { return [] } + if batchSize == 1 { + return [try forwardOptimized(tokenId: tokenIds[0], position: positions[0])] + } + + let cmdBuf = engine.commandQueue.makeCommandBuffer()! + let device = engine.device + + let batchInputBuffer = device.makeBuffer( + length: batchSize * hiddenSize * MemoryLayout.stride + )! + let batchOutputBuffer = device.makeBuffer( + length: batchSize * vocabSize * MemoryLayout.stride + )! + + // Process embeddings in batch + var batchEmbeddings: [[Float]] = [] + for i in 0..( + start: outputPtr + offset, + count: vocabSize + )) + results.append(logits) + } + + return results + } + + /// Helper: Dequantize embedding for a single token + private func dequantizeEmbedding(tokenId: Int) throws -> [Float] { + let device = engine.device + let tempBuffer = device.makeBuffer(length: hiddenSize * 4)! + let cmdBuf = engine.commandQueue.makeCommandBuffer()! + + try dequantizeRowOptimized( + weight: embedWeight, + tokenId: tokenId, + output: tempBuffer, + cmdBuf: cmdBuf + ) + + if embedScale != 1.0 { + try scaleBufferOptimized(tempBuffer, scale: embedScale, count: hiddenSize, cmdBuf: cmdBuf) + } + + cmdBuf.commit() + cmdBuf.waitUntilCompleted() + + let ptr = tempBuffer.contents().assumingMemoryBound(to: Float.self) + return Array(UnsafeBufferPointer(start: ptr, count: hiddenSize)) + } + + /// Helper: Process layers for batch generation + private func processLayersBatch( + input: MTLBuffer, + position: Int, + cmdBuf: MTLCommandBuffer, + outputOffset: Int + ) throws { + // For now, use existing layer forward (batching would require kernel modification) + // This still saves embedding/lm_head dispatches + + let h = input + + // Process through all layers + for layerIdx in 0.. 0 ? + layerIdx * perLayerInputSize * MemoryLayout.stride : 0 + + try layers[layerIdx].forwardOptimized( + input: h, + position: position, + kvCache: cache, + shouldStoreKV: isOwner, + temps: temps, + engine: engine, + cmdBuf: cmdBuf, + perLayerInput: perLayerEmbedBuffer, + perLayerInputOffset: plOffset + ) + } + + // Final norm + var lmInput = h + if let fn = finalNorm { + try rmsNormOptimized(input: h, weight: fn, output: temps.ns, + count: hiddenSize, cmdBuf: cmdBuf) + lmInput = temps.ns + } + + // LM head (batched output) + // Note: This would need special handling for true batching + try quantizedMatmulOptimized( + input: lmInput, + weights: embedWeight, + output: logitsBuffer, + cmdBuf: cmdBuf + ) + + // Logits scaling + if embedWeight.groupSize == 32 && embedWeight.inDim == hiddenSize { + let logitsScale = Float(30.0 / 116.23 / sqrt(Float(hiddenSize))) + try scaleBufferOptimized(logitsBuffer, scale: logitsScale, count: vocabSize, cmdBuf: cmdBuf) + } + + // Softcapping + if let cap = finalLogitSoftcapping { + try applyLogitSoftcappingOptimized( + buffer: logitsBuffer, + cap: cap, + count: vocabSize, + cmdBuf: cmdBuf + ) + } + } + + /// Batch generation - generate N tokens with batch processing + public func generateBatch(startToken: Int, numTokens: Int) throws -> [Int] { + var tokens: [Int] = [startToken] + var allLogits: [[Float]] = [] + + // Generate in batches of 8 (optimal for most GPUs) + let batchSize = 8 + + while tokens.count < numTokens { + let remaining = numTokens - tokens.count + let currentBatchSize = min(batchSize, remaining) + + // Prepare batch inputs + var batchTokens: [Int] = [] + var batchPositions: [Int] = [] + + for i in 0..= numTokens { break } + } + } + + return tokens + } + + /// Helper: Argmax for token selection + private func argmax(_ logits: [Float]) -> Int { + var maxIdx = 0 + var maxVal = logits[0] + for i in 1.. maxVal { + maxVal = logits[i] + maxIdx = i + } + } + return maxIdx + } +} \ No newline at end of file diff --git a/Sources/MarkBase/BatchGenerationOptimized.swift b/Sources/MarkBase/BatchGenerationOptimized.swift new file mode 100644 index 0000000..f6f3f34 --- /dev/null +++ b/Sources/MarkBase/BatchGenerationOptimized.swift @@ -0,0 +1,226 @@ +import Metal + +// Production-ready batch generation with buffer reuse +// Eliminates buffer allocation overhead for maximum performance + +extension E4BModel { + + /// Batch inference context - reuses buffers across calls + public final class BatchContext { + let device: MTLDevice + let maxBatchSize: Int + let hiddenSize: Int + let vocabSize: Int + + // Reusable buffers + let batchInputBuffer: MTLBuffer + let batchOutputBuffer: MTLBuffer + let tempEmbeddingBuffer: MTLBuffer + + public init(device: MTLDevice, maxBatchSize: Int, hiddenSize: Int, vocabSize: Int) { + self.device = device + self.maxBatchSize = maxBatchSize + self.hiddenSize = hiddenSize + self.vocabSize = vocabSize + + // Pre-allocate buffers + self.batchInputBuffer = device.makeBuffer( + length: maxBatchSize * hiddenSize * MemoryLayout.stride, + options: .storageModeShared + )! + self.batchOutputBuffer = device.makeBuffer( + length: maxBatchSize * vocabSize * MemoryLayout.stride, + options: .storageModeShared + )! + self.tempEmbeddingBuffer = device.makeBuffer( + length: hiddenSize * MemoryLayout.stride, + options: .storageModeShared + )! + } + } + + /// Create a batch context for reuse (call once at startup) + public func createBatchContext(maxBatchSize: Int = 8) -> BatchContext { + return BatchContext( + device: engine.device, + maxBatchSize: maxBatchSize, + hiddenSize: hiddenSize, + vocabSize: vocabSize + ) + } + + /// Optimized batch forward with buffer reuse + public func forwardBatchOptimized( + tokenIds: [Int], + positions: [Int], + context: BatchContext + ) throws -> [[Float]] { + guard tokenIds.count == positions.count else { + throw NSError(domain: "Batch", code: -1, + userInfo: [NSLocalizedDescriptionKey: "tokenIds and positions must match"]) + } + + let batchSize = tokenIds.count + guard batchSize <= context.maxBatchSize else { + throw NSError(domain: "Batch", code: -2, + userInfo: [NSLocalizedDescriptionKey: "Batch size exceeds context max"]) + } + + if batchSize == 0 { return [] } + if batchSize == 1 { + return [try forwardOptimized(tokenId: tokenIds[0], position: positions[0])] + } + + // ── Phase 1: Process embeddings SEPARATELY (must complete before layers) ── + let embedCmdBuf = engine.commandQueue.makeCommandBuffer()! + let inputPtr = context.batchInputBuffer.contents().assumingMemoryBound(to: Float.self) + + for i in 0..