Initial commit: E4B-MarkBase model integration with passing tests
CI / build-and-test (push) Has been cancelled

- E4B-MarkBase model (42 layers, 4.4GB) loaded successfully
- All Phase 1-6 tests passed (model loading, forward pass, vision/audio towers, token generation, performance)
- All stress tests passed (5/5 in 127.6s)
  - Concurrent inference
  - Memory stress (67.5 tok/s, 0 NaN)
  - Continuous generation
  - Batch processing
  - Long-running stability
- Swift Metal inference engine with multimodal support
This commit is contained in:
MarkBase Admin
2026-06-23 18:12:35 +08:00
commit ac75faa0cc
301 changed files with 63426 additions and 0 deletions
+51
View File
@@ -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
}
}
@@ -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..<numFrames {
let startIdx = frameIdx * hopLength
let endIdx = min(startIdx + nFft, audioData.count)
// Zero-pad if frame is shorter than nFft
var frame = [Float](repeating: 0, count: nFft)
for i in startIdx..<endIdx {
frame[i - startIdx] = audioData[i]
}
let windowedFrame = applyHannWindow(frame)
let spectrum = computeSpectrum(windowedFrame)
let melEnergies = computeMelEnergies(spectrum)
melSpec[frameIdx] = melEnergies
}
return melSpec
}
private func applyHannWindow(_ frame: [Float]) -> [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..<spectrum.count {
var real: Float = 0
var imag: Float = 0
for i in 0..<n {
let angle = -2.0 * Float.pi * Float(k) * Float(i) / Float(n)
real += frame[i] * cos(angle)
imag += frame[i] * sin(angle)
}
spectrum[k] = sqrt(real * real + imag * imag)
}
return spectrum
}
private func computeMelEnergies(_ spectrum: [Float]) -> [Float] {
var melEnergies = [Float](repeating: 0, count: nMels)
let melPoints = createMelFilterbank()
for melIdx in 0..<nMels {
var sum: Float = 0
for fftIdx in 0..<spectrum.count {
sum += spectrum[fftIdx] * melPoints[melIdx][fftIdx]
}
melEnergies[melIdx] = log10(max(sum, 1e-10))
}
return melEnergies
}
private func createMelFilterbank() -> [[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..<nMels + 2).map { i in
melMin + Float(i) * (melMax - melMin) / Float(nMels + 1)
}
let hzPoints = melPoints.map { melToHz($0) }
let binPoints = hzPoints.map { Int(round($0 * Float(nFft) / Float(sampleRate))) }
for melIdx in 0..<nMels {
let leftBin = binPoints[melIdx]
let centerBin = binPoints[melIdx + 1]
let rightBin = binPoints[melIdx + 2]
for bin in leftBin..<centerBin {
if bin < filterbank[melIdx].count {
filterbank[melIdx][bin] = Float(bin - leftBin) / Float(centerBin - leftBin)
}
}
for bin in centerBin..<rightBin {
if bin < filterbank[melIdx].count {
filterbank[melIdx][bin] = Float(rightBin - bin) / Float(rightBin - centerBin)
}
}
}
return filterbank
}
private func hzToMel(_ hz: Float) -> 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..<nMels {
for b in 0..<spectrumSize {
flatFilterbank[m * spectrumSize + b] = filterbank2D[m][b]
}
}
let cmdBuf = engine.commandQueue.makeCommandBuffer()!
let audioBuf = device.makeBuffer(bytes: audioData, length: audioData.count * 4)!
let filterbankBuf = device.makeBuffer(bytes: flatFilterbank, length: flatFilterbank.count * 4)!
let spectrumBuf = device.makeBuffer(length: numFrames * spectrumSize * 4)!
let melBuf = device.makeBuffer(length: melBufferSize * 4)!
let psoDFT = try engine.pipeline(named: "audio_dft_magnitude")
let psoMel = try engine.pipeline(named: "audio_mel_filterbank")
do {
let enc = cmdBuf.makeComputeCommandEncoder()!
enc.setComputePipelineState(psoDFT)
enc.setBuffer(audioBuf, offset: 0, index: 0)
enc.setBuffer(spectrumBuf, offset: 0, index: 1)
var n = UInt32(nFft); enc.setBytes(&n, length: 4, index: 2)
var h = UInt32(hopLength); enc.setBytes(&h, length: 4, index: 3)
var nf = UInt32(numFrames); enc.setBytes(&nf, length: 4, index: 4)
var ss = UInt32(spectrumSize); enc.setBytes(&ss, length: 4, index: 5)
var al = UInt32(audioData.count); enc.setBytes(&al, length: 4, index: 6)
let grid = MTLSize(width: numFrames, height: spectrumSize, depth: 1)
let tg = MTLSize(width: 8, height: 8, depth: 1)
enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
enc.endEncoding()
}
do {
let enc = cmdBuf.makeComputeCommandEncoder()!
enc.setComputePipelineState(psoMel)
enc.setBuffer(spectrumBuf, offset: 0, index: 0)
enc.setBuffer(filterbankBuf, offset: 0, index: 1)
enc.setBuffer(melBuf, offset: 0, index: 2)
var ss = UInt32(spectrumSize); enc.setBytes(&ss, length: 4, index: 3)
var nm = UInt32(nMels); enc.setBytes(&nm, length: 4, index: 4)
var nf = UInt32(numFrames); enc.setBytes(&nf, length: 4, index: 5)
let grid = MTLSize(width: numFrames, height: nMels, depth: 1)
let tg = MTLSize(width: 8, height: 8, depth: 1)
enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
enc.endEncoding()
}
cmdBuf.commit()
cmdBuf.waitUntilCompleted()
let ptr = melBuf.contents().assumingMemoryBound(to: Float.self)
let flat = Array(UnsafeBufferPointer(start: ptr, count: melBufferSize))
var result = [[Float]](repeating: [Float](repeating: 0, count: nMels), count: numFrames)
for f in 0..<numFrames {
for m in 0..<nMels {
result[f][m] = flat[f * nMels + m]
}
}
return result
}
public func loadAudioFile(url: URL) throws -> [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<Float>.stride)
CMBlockBufferCopyDataBytes(blockBuffer, atOffset: 0, dataLength: length, destination: &data)
samples.append(contentsOf: data)
}
}
}
return samples
}
}
+740
View File
@@ -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
)
}
}
+129
View File
@@ -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)
}
}
+609
View File
@@ -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<Float>.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..<config.numHiddenLayers {
loadedLayers.append(try AudioLayerWeightsE2B(device: device, layerIdx: i, floats: floats))
}
layers = loadedLayers
}
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<Float>.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)
}
+209
View File
@@ -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..<config.numHiddenLayers {
loadedLayers.append(try AudioLayerWeights(device: device, layerIdx: i,
tensors: tensors, floats: floats,
descriptors: descriptors))
}
layers = loadedLayers
}
// Helpers
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<Float>.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<Float>.stride,
options: .storageModeShared) else {
throw WeightError.bufferCreationFailed(sName)
}
guard let bBuf = device.makeBuffer(bytes: bFloats, length: bFloats.count * MemoryLayout<Float>.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<Float>.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<Float>.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")
}
}