SpeechKit/Sources/ASRKit/LMDecoder.swift

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)
                }
            }
        }
    }
}