File size: 8,729 Bytes
64f7370 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 | import Foundation
import SpeechCore
import OnnxRuntimeBindings
/// int4 LM decoder (ORT CPU): prefill + per-token decode with external KV cache,
/// plus the tiny mask_gen graph for the audio tower's attention/valid masks.
package final class LMDecoder {
package static let hidden = 512
package static let numLayers = 8
package static let numKVHeads = 8
package static let headDim = 64
package static let maxTotalLen = 512
private let env: ORTEnv
private let prefill: ORTSession
private let decode: ORTSession
private let maskGen: ORTSession
private let prefillOutputs: [String]
private let decodeOutputs: [String]
private let store: AssetStore
// KV cache: persistent NSMutableData-backed ORTValues, mutated in place
// between steps (avoids re-allocating/copying 16 MB per generated token).
private var cacheData: [NSMutableData] = []
private var cacheValues: [ORTValue] = []
private func resetCaches() throws {
let bytes = Self.numKVHeads * Self.maxTotalLen * Self.headDim * MemoryLayout<Float>.size
if cacheData.isEmpty {
for _ in 0..<(Self.numLayers * 2) {
let d = NSMutableData(length: bytes)!
cacheData.append(d)
cacheValues.append(try ORTValue(
tensorData: d, elementType: .float,
shape: [1, Self.numKVHeads, Self.maxTotalLen, Self.headDim].map { NSNumber(value: $0) }))
}
} else {
for d in cacheData { memset(d.mutableBytes, 0, d.length) }
}
}
package init(prefillURL: URL, decodeURL: URL, maskGenURL: URL, store: AssetStore) throws {
env = try ORTEnv(loggingLevel: .warning)
let opts = try ORTSessionOptions()
try opts.setIntraOpNumThreads(2)
prefill = try ORTSession(env: env, modelPath: prefillURL.path, sessionOptions: opts)
decode = try ORTSession(env: env, modelPath: decodeURL.path, sessionOptions: opts)
maskGen = try ORTSession(env: env, modelPath: maskGenURL.path, sessionOptions: opts)
prefillOutputs = try prefill.outputNames()
decodeOutputs = try decode.outputNames()
self.store = store
}
// MARK: - tensors
private static func tensor(_ values: [Float], shape: [Int]) throws -> ORTValue {
let data = NSMutableData(bytes: values, length: values.count * 4)
return try ORTValue(tensorData: data, elementType: .float,
shape: shape.map { NSNumber(value: $0) })
}
private static func tensor(_ values: [Int64], shape: [Int]) throws -> ORTValue {
let data = NSMutableData(bytes: values, length: values.count * 8)
return try ORTValue(tensorData: data, elementType: .int64,
shape: shape.map { NSNumber(value: $0) })
}
private static func floats(_ value: ORTValue) throws -> [Float] {
let data = try value.tensorData() as Data
return data.withUnsafeBytes { Array($0.bindMemory(to: Float.self)) }
}
// MARK: - mask generation for the audio tower
/// Returns (attnMask 1*1*390*390, validMask 390)
package func generateMasks(encLen: Int) throws -> (attn: [Float], valid: [Bool]) {
let out = try maskGen.run(
withInputs: ["audio_feature_lengths": try Self.tensor([Int64(encLen)], shape: [1])],
outputNames: Set(try maskGen.outputNames()), runOptions: nil)
let names = try maskGen.outputNames()
let attn = try Self.floats(out[names[0]]!)
let validData = try out[names[1]]!.tensorData() as Data
let valid = validData.withUnsafeBytes {
$0.bindMemory(to: Int64.self).map { $0 != 0 }
}
return (attn, valid)
}
// MARK: - generation
package struct GenerationResult {
package let tokenIDs: [Int]
package let text: String
package let hitStop: Bool
}
package func generate(promptEmbeds: [[Float]], promptIDs: [Int],
maxNewTokens: Int = 128,
isCancelled: () -> Bool = { false }) throws -> GenerationResult {
let promptLen = promptEmbeds.count
guard promptLen < Self.maxTotalLen - 1 else {
throw SpeechError.promptTooLong(tokens: promptLen, limit: Self.maxTotalLen - 1)
}
// leave room in the fixed-size cache for generation
let budget = min(maxNewTokens, Self.maxTotalLen - promptLen)
try resetCaches()
// prefill
var logits: [Float] = try autoreleasepool {
let flat = promptEmbeds.flatMap { $0 }
let feeds: [String: ORTValue] = [
"inputs_embeds": try Self.tensor(flat, shape: [1, promptLen, Self.hidden]),
"cache_position": try Self.tensor((0..<promptLen).map(Int64.init), shape: [promptLen]),
]
let pOut = try prefill.run(withInputs: feeds, outputNames: Set(prefillOutputs), runOptions: nil)
for layer in 0..<Self.numLayers {
try writeCache(layer: layer, keyValue: pOut, names: prefillOutputs,
position: 0, length: promptLen)
}
return try lastLogits(from: pOut[prefillOutputs[0]]!)
}
// decode loop
let blocked = Set(store.manifest.tokens.extra_block_token_ids)
let eos = Set(store.manifest.tokens.eos_token_ids)
let pad = store.manifest.tokens.pad_token_id
var generated: [Int] = []
var hitStop = false
var position = promptLen
for _ in 0..<budget {
if isCancelled() { throw SpeechError.cancelled }
for b in blocked where b < logits.count { logits[b] = -.infinity }
var next = 0
var best = -Float.infinity
for (i, v) in logits.enumerated() where v > best { best = v; next = i }
if eos.contains(next) || next == pad { hitStop = true; break }
generated.append(next)
logits = try autoreleasepool {
var mask = [Int64](repeating: 0, count: Self.maxTotalLen)
for i in 0...position { mask[i] = 1 }
var feeds: [String: ORTValue] = [
"inputs_embeds": try Self.tensor(store.embedding(for: next), shape: [1, 1, Self.hidden]),
"attention_mask": try Self.tensor(mask, shape: [1, Self.maxTotalLen]),
"cache_position": try Self.tensor([Int64(position)], shape: [1]),
]
for layer in 0..<Self.numLayers {
feeds["cache_key_\(layer)"] = cacheValues[layer * 2]
feeds["cache_value_\(layer)"] = cacheValues[layer * 2 + 1]
}
let dOut = try decode.run(withInputs: feeds, outputNames: Set(decodeOutputs), runOptions: nil)
for layer in 0..<Self.numLayers {
try writeCache(layer: layer, keyValue: dOut, names: decodeOutputs,
position: position, length: 1)
}
return try lastLogits(from: dOut[decodeOutputs[0]]!)
}
position += 1
}
return GenerationResult(tokenIDs: generated, text: store.decode(generated), hitStop: hitStop)
}
private func lastLogits(from value: ORTValue) throws -> [Float] {
let all = try Self.floats(value)
let info = try value.tensorTypeAndShapeInfo()
let vocab = Int(truncating: info.shape.last!)
return Array(all[(all.count - vocab)...])
}
/// Copy per-layer (1,8,L,64) new keys/values into the persistent cache
/// buffers at `position` (direct memcpy, no intermediate arrays).
private func writeCache(layer: Int, keyValue: [String: ORTValue],
names: [String], position: Int, length: Int) throws {
let keyData = try keyValue[names[1 + 2 * layer]]!.tensorData() as Data
let valData = try keyValue[names[2 + 2 * layer]]!.tensorData() as Data
let rowBytes = Self.headDim * MemoryLayout<Float>.size
let strideCacheBytes = Self.maxTotalLen * rowBytes
let strideNewBytes = length * rowBytes
for (data, cache) in [(keyData, cacheData[layer * 2]), (valData, cacheData[layer * 2 + 1])] {
data.withUnsafeBytes { (src: UnsafeRawBufferPointer) in
let dstBase = cache.mutableBytes
for h in 0..<Self.numKVHeads {
memcpy(dstBase + h * strideCacheBytes + position * rowBytes,
src.baseAddress! + h * strideNewBytes,
strideNewBytes)
}
}
}
}
}
|