Potential new approach to quantized [KVCache]

TokenIterator.swift

//  TokenIterator.swift

import Foundation
import MLX
import MLXLMCommon

/// Generator of tokens.
///
/// This is typically used via a call to ``generate(input:cache:parameters:context:wiredMemoryTicket:tools:)`` returning `AsyncStream<Generation>`.
///
/// To use it directly:
///
/// ```swift
/// let generateParameters: GenerateParameters
/// let input: LMInput
/// let model: LanguageModel
///
/// let iterator = try TokenIterator(input: input, model: model, parameters: generateParameters)
///
/// for token in iterator {
///     ...
/// }
/// ```
///
/// Tokens are integers that can be passed through a `Tokenizer` or ``StreamingDetokenizer`` to produce Strings.
///
/// Port of `generate_step()` from https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/utils.py
///
/// Note: this uses `asyncEval()` and there may be an async evaluation running after a call to `next()`.
public struct TokenIterator: TokenIteratorProtocol {
    let model: any LanguageModel
    var state: LMOutput.State?

    var y: LMInput.Text
    var cacheContainer: KVCacheContainer
    var cache: [KVCache] { cacheContainer.cache }
    var processor: LogitProcessor?
    let sampler: LogitSampler

    private(set) public var tokenCount = 0
    public let maxTokens: Int?

    // Cache quantization parameters
    let kvBits: Int?
    let kvGroupSize: Int
    let quantizedKVStart: Int

    // Internal metrics
    private(set) public var promptPrefillTime: TimeInterval = 0.0

    /// Initialize a `TokenIterator` with the given input.
    ///
    /// - Parameters:
    ///   - input: language model input
    ///   - model: the ``LanguageModel``
    ///   - cache: optional ``KVCache``
    ///   - parameters: the generation parameters
    ///   - processor: the logit processor override
    ///   - sampler: the logit sampler override
    public init(input: LMInput,
                model: any LanguageModel,
                cacheContainer: KVCacheContainer? = nil,
                parameters: GenerateParameters,
                processor: LogitProcessor? = nil,
                sampler: LogitSampler? = nil) throws {
        self.model = model
        self.y = input.text
        self.cacheContainer = cacheContainer ?? KVCacheContainer(cache: nil, model: model, parameters: parameters)

        self.processor = processor ?? parameters.processor()
        self.sampler = sampler ?? parameters.sampler()
        self.maxTokens = parameters.maxTokens

        self.kvBits = parameters.kvBits
        self.kvGroupSize = parameters.kvGroupSize
        self.quantizedKVStart = parameters.quantizedKVStart

        let start = Date.timeIntervalSinceReferenceDate
        try prepare(input: input, windowSize: parameters.prefillStepSize)
        self.promptPrefillTime = Date.timeIntervalSinceReferenceDate - start
    }

    /// Initialize a `TokenIterator` with the given tokens. Note: this has been
    /// replaced with ``init(input:model:cache:parameters:)``.
    ///
    /// - Parameters:
    ///   - prompt: the prompt tokens
    ///   - model: the ``LanguageModel``
    ///   - cache: optional ``KVCache``
    ///   - parameters: the generation parameters
    @available(*, deprecated, message: "please use init(input:model:cache:parameters:)")
    public init(
        prompt: MLXArray, model: any LanguageModel, cache: [KVCache]? = nil,
        parameters: GenerateParameters
    ) throws {
        let input = LMInput(tokens: prompt)
        try self.init(input: input, model: model, cache: cache, parameters: parameters)
    }
    
    /// Initialize a `TokenIterator` with the given input.
    ///
    /// - Parameters:
    ///   - input: language model input
    ///   - model: the ``LanguageModel``
    ///   - cache: optional ``KVCache``
    ///   - parameters: the generation parameters
    ///   - processor: the logit processor override
    ///   - sampler: the logit sampler override
    public init(input: LMInput,
                model: any LanguageModel,
                cache: [KVCache]? = nil,
                parameters: GenerateParameters,
                processor: LogitProcessor? = nil,
                sampler: LogitSampler? = nil) throws {
        let container = KVCacheContainer(cache: cache, model: model, parameters: parameters)

        try self.init(input: input,
                      model: model,
                      cacheContainer: container,
                      parameters: parameters,
                      processor: processor,
                      sampler: sampler)
    }

    /// Initialize a `TokenIterator` with the given input and logit handling.
    ///
    /// - Parameters:
    ///   - input: language model input
    ///   - model: the ``LanguageModel``
    ///   - cache: optional ``KVCache``
    ///   - processor: the logit processor
    ///   - sampler: the logit sampler
    ///   - prefillStepSize: optional prefill step size
    ///   - maxTokens: maximum number of tokens to generate
    public init(
        input: LMInput, model: any LanguageModel, cache: [KVCache]? = nil,
        processor: LogitProcessor?, sampler: LogitSampler, prefillStepSize: Int = 512,
        maxTokens: Int? = nil
    ) throws {
        let parameters = GenerateParameters(maxTokens: maxTokens,
                                            kvBits: nil,
                                            kvGroupSize: 64,
                                            quantizedKVStart: 0,
                                            prefillStepSize: prefillStepSize)

        try self.init(input: input,
                      model: model,
                      cache: cache,
                      parameters: parameters,
                      processor: processor,
                      sampler: sampler)
    }

    mutating func prepare(input: LMInput, windowSize: Int? = nil) throws {
        processor?.prompt(input.text.tokens)

        switch try model.prepare(input, cache: cache, windowSize: windowSize) {
        case .tokens(let tokens):
            y = tokens

            // evaluate the remainder of the prompt -- this primes the pump
            let token = step(previous: y)
            y = .init(tokens: token)
            asyncEval(y.tokens)

        case .logits(let result):
            y = .init(tokens: convertToToken(logits: result.logits))
            asyncEval(y.tokens)
        }
    }

    mutating func convertToToken(logits: MLXArray) -> MLXArray {
        // process the logits (one hot array of possible tokens)
        var logits = logits[0..., -1, 0...]
        logits = processor?.process(logits: logits) ?? logits

        // transform logits back to a token
        let y = sampler.sample(logits: logits)

        processor?.didSample(token: y)

        return y
    }

    /// Evaluate the next token and return the new token (y), updating cache state
    mutating func step(previous: LMInput.Text) -> MLXArray {
        let result = model(previous[text: .newAxis], cache: cache.isEmpty ? nil : cache, state: state)
        self.state = result.state

        // Attempt dynamic cache quantization after each step if needed
        cacheContainer.maybeQuantize(kvBits: kvBits, kvGroupSize: kvGroupSize, quantizedKVStart: quantizedKVStart)

        return convertToToken(logits: result.logits)
    }

    mutating public func next() -> Int? {
        if let maxTokens, tokenCount >= maxTokens {
            return nil
        }

        // save current value -- this will be returned
        let previousY = y

        // compute the next state and async eval the next token
        let token = step(previous: previousY)
        y = .init(tokens: token)
        asyncEval(token)

        tokenCount += 1

        return previousY.tokens.item(Int.self)
    }
}

public protocol QuantizableKVCache: KVCache {
    func toQuantized(groupSize: Int, bits: Int) -> QuantizedKVCache
}

extension KVCacheSimple: QuantizableKVCache { }

extension KVCache {
    var keyHeadDim: Int? { state.count == 2 ? state[0].dim(3) : nil }
    var valueHeadDim: Int? { state.count == 2 ? state[1].dim(3) : nil }

    func asQuantizableWithGroupSize(_ groupSize: Int) -> QuantizableKVCache? {
        switch self {
        case let simple as KVCacheSimple:
            return simple
        // TODO: RotatingKVCache.toQuantized() is not implemented yet, like in Python.
        // When implemented, add: case let rotating as RotatingKVCache...
        default:
            return nil
        }
    }

    func canQuantize(at quantizedKVStart: Int, groupSize: Int) -> Bool {
        guard let quantizable = asQuantizableWithGroupSize(groupSize) else { return false }

        return quantizable.offset >= quantizedKVStart
    }

    func isValidQuantizedGroupSize(_ groupSize: Int) -> Bool {
        guard let keyHeadDim, let valueHeadDim else { return false }

        return [32, 64, 128].contains { keyHeadDim.isMultiple(of: $0) && valueHeadDim.isMultiple(of: $0) }
    }

    func resolvedKVQuantizationGroupSize(_ groupSize: Int) -> Int? {
        guard let keyHeadDim, let valueHeadDim else { return nil }
        let compatible = [32, 64, 128].filter {
            keyHeadDim.isMultiple(of: $0) && valueHeadDim.isMultiple(of: $0)
        }
        guard !compatible.isEmpty else { return nil }
        let requested = max(1, groupSize)

        return compatible.min { lhs, rhs in
            let lhsDistance = abs(lhs - requested)
            let rhsDistance = abs(rhs - requested)
            if lhsDistance == rhsDistance {
                return lhs < rhs
            }
            return lhsDistance < rhsDistance
        }
    }
}

public class KVCacheContainer {
    private(set) public var cache: [KVCache] = []
    private var canQuantize: Bool

    convenience public init(cache: [KVCache]? = nil, model: any LanguageModel, parameters: GenerateParameters) {
        let cache = cache ?? model.newCache(parameters: parameters)

        self.init(cache: cache, kvBits: parameters.kvBits, kvGroupSize: parameters.kvGroupSize)
    }

    public init(cache: [KVCache], kvBits: Int?, kvGroupSize: Int) {
        self.cache = cache
        self.canQuantize = kvBits != nil &&
            cache.contains { $0.asQuantizableWithGroupSize(kvGroupSize) != nil }
    }

    public func maybeQuantize(kvBits: Int?, kvGroupSize: Int = 64, quantizedKVStart: Int = 0) {
        guard canQuantize else { return }

        canQuantize = !maybeQuantizeKVCache2(cache: &cache, kvBits: kvBits, kvGroupSize: kvGroupSize, quantizedKVStart: quantizedKVStart)
    }
}

@discardableResult
public func maybeQuantizeKVCache2(cache: inout [KVCache],
                                  kvBits: Int?,
                                  kvGroupSize: Int = 64,
                                  quantizedKVStart: Int = 0) -> Bool {
    guard let kvBits = kvBits, !cache.isEmpty else { return false }
    guard cache.contains(where: { $0.canQuantize(at: quantizedKVStart, groupSize: kvGroupSize) }) else { return false }

    cache = cache.map {
        guard let quantizableCache = $0.asQuantizableWithGroupSize(kvGroupSize) else { return $0 }
        guard quantizableCache.isValidQuantizedGroupSize(kvGroupSize) else { return $0 }

        return quantizableCache.toQuantized(groupSize: kvGroupSize, bits: kvBits)
    }

    return true
}