Merge branch 'main' into sukru/state-handlers

Author: stikves

python/src/coreai_models/diffusion/pipeline.py

@@ -162,8 +162,6 @@ def _load_hf_pipeline(model_id: str, pipeline_type: str, model_dtype: torch.dtyp
         from diffusers import Flux2KleinPipeline
 
         hf_pipe = Flux2KleinPipeline.from_pretrained(model_id, torch_dtype=model_dtype)
-        # Text encoder needs float32 for token embedding precision
-        hf_pipe.text_encoder = hf_pipe.text_encoder.float()
         return hf_pipe
 
     if pipeline_type == "sd3":

swift/Sources/CoreAIDiffusionPipeline/Components/CoreAIDenoiser.swift

@@ -31,7 +31,7 @@ public final class CoreAIDenoiser: Sendable {
         let batchSize = latents.shape[0]
 
         var timestepArray = NDArray(shape: [batchSize], scalarType: .float32)
-        var timestepView = timestepArray.mutableView(as: Float.self)
+        let timestepView = timestepArray.mutableView(as: Float.self)
         timestepView.withUnsafeMutablePointer { ptr, _, _ in
             for i in 0..<batchSize { ptr[i] = timestep }
         }
@@ -58,7 +58,7 @@ public final class CoreAIDenoiser: Sendable {
 
         let shape = latents.shape
         var result = NDArray(shape: shape, scalarType: .float32)
-        var resultView = result.mutableView(as: Float.self)
+        let resultView = result.mutableView(as: Float.self)
         resultView.withUnsafeMutablePointer { ptr, _, _ in
             for i in 0..<floats.count { ptr[i] = floats[i] }
         }

swift/Sources/CoreAIDiffusionPipeline/Components/CoreAIDiffusionModelFunction.swift

@@ -55,13 +55,13 @@ public actor CoreAIDiffusionModelFunction {
             switch resolved.scalarType {
             #if !((os(macOS) || targetEnvironment(macCatalyst)) && arch(x86_64))
             case .float16:
-                var view = array.mutableView(as: Float16.self)
+                let view = array.mutableView(as: Float16.self)
                 view.withUnsafeMutablePointer { ptr, _, _ in
                     for j in 0..<data.count { ptr[j] = Float16(data[j]) }
                 }
             #endif
             case .float32:
-                var view = array.mutableView(as: Float.self)
+                let view = array.mutableView(as: Float.self)
                 view.withUnsafeMutablePointer { ptr, _, _ in
                     for j in 0..<data.count { ptr[j] = data[j] }
                 }
@@ -83,7 +83,7 @@ public actor CoreAIDiffusionModelFunction {
             guard case .ndArray(let nd) = fn.descriptor.inputDescriptor(of: name) else { continue }
             let resolved = nd.resolvingDynamicDimensions(shape)
             var array = NDArray(descriptor: resolved)
-            var view = array.mutableView(as: Int32.self)
+            let view = array.mutableView(as: Int32.self)
             view.withUnsafeMutablePointer { ptr, _, _ in
                 for j in 0..<data.count { ptr[j] = data[j] }
             }

swift/Sources/CoreAIDiffusionPipeline/Components/CoreAILatentCodec.swift

@@ -44,7 +44,7 @@ public final class CoreAILatentDecoder: Sendable {
         let outW = shape[3] * 8
         let outShape = [1, 3, outH, outW]
         var result = NDArray(shape: outShape, scalarType: .float32)
-        var resultView = result.mutableView(as: Float.self)
+        let resultView = result.mutableView(as: Float.self)
         resultView.withUnsafeMutablePointer { ptr, _, _ in
             for i in 0..<outputFloats.count { ptr[i] = outputFloats[i] }
         }
@@ -107,7 +107,7 @@ public final class CoreAILatentEncoder: Sendable {
         // Scale and return as NDArray
         let outShape = [1, 4, height / 8, width / 8]
         var result = NDArray(shape: outShape, scalarType: .float32)
-        var resultView = result.mutableView(as: Float.self)
+        let resultView = result.mutableView(as: Float.self)
         resultView.withUnsafeMutablePointer { ptr, _, _ in
             for i in 0..<outputFloats.count { ptr[i] = outputFloats[i] * scaleFactor }
         }

swift/Sources/CoreAIDiffusionPipeline/Pipelines/Flux2Pipeline.swift

@@ -274,7 +274,7 @@ public struct Flux2Pipeline: DiffusionPipeline {
         let image = try DiffusionUtilities.pixelsToCGImage(pixels, height: outputHeight, width: outputWidth)
 
         var latentsND = NDArray(shape: latentShape, scalarType: .float32)
-        var latentsView = latentsND.mutableView(as: Float.self)
+        let latentsView = latentsND.mutableView(as: Float.self)
         latentsView.withUnsafeMutablePointer { ptr, _, _ in
             for i in 0..<noise.count { ptr[i] = noise[i] }
         }

swift/Sources/CoreAIDiffusionPipeline/Pipelines/SD3Pipeline.swift

@@ -160,7 +160,7 @@ public struct SD3Pipeline: DiffusionPipeline {
             pixels, height: imageSize, width: imageSize)
 
         var latentsND = NDArray(shape: latentShape, scalarType: .float32)
-        var latentsView = latentsND.mutableView(as: Float.self)
+        let latentsView = latentsND.mutableView(as: Float.self)
         latentsView.withUnsafeMutablePointer { ptr, _, _ in
             for i in 0..<latents.count { ptr[i] = latents[i] }
         }
@@ -226,7 +226,7 @@ public struct SD3Pipeline: DiffusionPipeline {
         }
 
         var inputArray = NDArray(shape: [1, Self.clipSeqLen], scalarType: .int32)
-        var view = inputArray.mutableView(as: Int32.self)
+        let view = inputArray.mutableView(as: Int32.self)
         view.withUnsafeMutablePointer { ptr, _, _ in
             for i in 0..<ids.count { ptr[i] = ids[i] }
         }

swift/Sources/CoreAIDiffusionPipeline/Pipelines/StableDiffusionPipeline.swift

@@ -161,7 +161,7 @@ public struct StableDiffusionPipeline: DiffusionPipeline {
 
         // Wrap latents back to NDArray for GenerationResult
         var latentsND = NDArray(shape: latentShape, scalarType: .float32)
-        var latentsView = latentsND.mutableView(as: Float.self)
+        let latentsView = latentsND.mutableView(as: Float.self)
         latentsView.withUnsafeMutablePointer { ptr, _, _ in
             for i in 0..<latents.count { ptr[i] = latents[i] }
         }

swift/Sources/CoreAILanguageModels/GuidedGeneration/ConstrainedGenerationSession.swift

@@ -17,6 +17,8 @@ import Tokenizers
 /// Each session is tied to a specific JSON schema and vocabulary. It tracks
 /// the generation state and produces token masks that enforce schema compliance.
 public struct ConstrainedGenerationSession: ~Copyable {
+    static let maxRollbackTokens = 64
+
     private let tokenizerInfo: TokenizerInfo
     private let compiler: GrammarCompiler
     private let compiledGrammar: CompiledGrammar
@@ -93,7 +95,7 @@ public struct ConstrainedGenerationSession: ~Copyable {
         self.tokenizerInfo = tokenizerInfo
         self.compiler = GrammarCompiler(tokenizerInfo: tokenizerInfo)
         self.compiledGrammar = try compiler.compileJSONSchema(jsonSchema)
-        self.matcher = GrammarMatcher(compiledGrammar: compiledGrammar)
+        self.matcher = GrammarMatcher(compiledGrammar: compiledGrammar, maxRollbackTokens: Self.maxRollbackTokens)
         self.vocabularySize = tokenizerInfo.vocabularySize
         self.bitmaskSize = (vocabularySize + 31) / 32
         self.bitmaskBuffer = Array(repeating: 0, count: bitmaskSize)
@@ -193,6 +195,50 @@ public struct ConstrainedGenerationSession: ~Copyable {
         matcher.reset()
         allTokensBlocked = false
     }
+
+    /// Rollback the grammar state by N tokens. Returns false if rollback failed
+    /// (e.g., exceeds maxRollbackTokens budget).
+    @discardableResult
+    public mutating func rollback(_ numTokens: Int = 1) -> Bool {
+        guard numTokens >= 0 else { return false }
+        return matcher.rollback(numTokens)
+    }
+
+    /// Find the longest deterministic string from the current grammar state.
+    /// Does not change the matcher state. Returns nil if no jump-forward is possible.
+    public func findJumpForwardString() -> String? {
+        matcher.findJumpForwardString()
+    }
+
+    /// Result of filling a bitmask for the next token.
+    public enum BitmaskResult: Equatable {
+        /// Grammar is terminated or all tokens are blocked — generation should stop.
+        case terminated
+        /// All tokens are allowed — no mask needed, generate unconstrained.
+        case unconstrained
+        /// Bitmask was written — apply it to constrain sampling.
+        case constrained
+    }
+
+    /// Fill the bitmask directly into a caller-provided buffer (e.g., a GPU-visible MTLBuffer).
+    ///
+    /// The caller must ensure the pointer has room for at least `(vocabularySize + 31) / 32`
+    /// Int32 words.
+    public mutating func fillBitmask(into pointer: UnsafeMutablePointer<Int32>) -> BitmaskResult {
+        if isTerminated { return .terminated }
+
+        let needsApplication = matcher.fillNextTokenBitmask(pointer)
+        if !needsApplication {
+            // xgrammar signals all tokens are allowed — no mask needed
+            return .unconstrained
+        }
+        // Check for all-zeros (no tokens allowed — grammar done)
+        for i in 0..<bitmaskSize {
+            if pointer[i] != 0 { return .constrained }
+        }
+        allTokensBlocked = true
+        return .terminated
+    }
 }
 
 // MARK: - Float16 Masking

swift/Sources/CoreAILanguageModels/GuidedGeneration/XGrammarWrapper.swift

@@ -136,9 +136,29 @@ public final class GrammarMatcher {
         return xgrammar_matcher_is_terminated(handle)
     }
 
+    public var isCompleted: Bool {
+        return xgrammar_matcher_is_completed(handle)
+    }
+
     public func reset() {
         xgrammar_matcher_reset(handle)
     }
+
+    @discardableResult
+    public func rollback(_ numTokens: Int = 1) -> Bool {
+        xgrammar_matcher_rollback(handle, Int32(numTokens))
+    }
+
+    /// Returns the longest deterministic string from the current grammar state,
+    /// or nil if no jump-forward is possible. Does not change matcher state.
+    public func findJumpForwardString() -> String? {
+        guard let cStr = xgrammar_matcher_find_jump_forward_string(handle) else {
+            return nil
+        }
+        let result = String(cString: cStr)
+        free(UnsafeMutablePointer(mutating: cStr))
+        return result.isEmpty ? nil : result
+    }
 }
 
 // MARK: - Errors

swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift

@@ -474,6 +474,85 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
         CLILogger.log("CoreAI clean engine cleanup complete")
         cleanupSpan.end()
     }
+
+    // MARK: - KV Cache (dynamic growth)
+
+    private func ensureKVCapacity(forContextLength needed: Int) throws {
+        guard needed > currentKVCapacity else { return }
+        guard needed <= config.maxContextLength else {
+            throw InferenceRuntimeError.invalidState(
+                "Context length \(needed) exceeds maximum \(config.maxContextLength)")
+        }
+
+        var newCapacity = currentKVCapacity
+        while newCapacity < needed { newCapacity *= 2 }
+        newCapacity = min(newCapacity, config.maxContextLength)
+
+        let resolvedKeyDesc = keyCacheDescriptor.resolvingDynamicDimensions(
+            keyCacheDescriptor.shape.map { $0 < 0 ? newCapacity : $0 })
+        let resolvedValueDesc = valueCacheDescriptor.resolvingDynamicDimensions(
+            valueCacheDescriptor.shape.map { $0 < 0 ? newCapacity : $0 })
+
+        var newKeyCache = NDArray(descriptor: resolvedKeyDesc)
+        var newValueCache = NDArray(descriptor: resolvedValueDesc)
+        _ = newKeyCache.mutableRawView()
+        _ = newValueCache.mutableRawView()
+
+        try Self.copyCache(from: keyCache, to: &newKeyCache)
+        try Self.copyCache(from: valueCache, to: &newValueCache)
+
+        CLILogger.log("KV cache grew: \(currentKVCapacity) → \(newCapacity)")
+        keyCache = newKeyCache
+        valueCache = newValueCache
+        currentKVCapacity = newCapacity
+    }
+
+    private static func copyCache(from source: NDArray, to destination: inout NDArray) throws {
+        let srcShape = source.shape
+        let dstShape = destination.shape
+        guard let headDim = srcShape.last else {
+            throw InferenceRuntimeError.invalidState("KV cache has empty shape — cannot copy")
+        }
+        let seqDim = KVCacheFactory.detectSequenceDim(shape: srcShape)
+
+        // Number of independent blocks before the sequence dimension (L * B * H or B * H)
+        let numBlocks = srcShape[..<seqDim].reduce(1, *)
+        let oldSeqLen = srcShape[seqDim]
+        let copySize = oldSeqLen * headDim
+
+        // Strides in elements for the sequence block
+        let srcBlockStride = srcShape[seqDim...].reduce(1, *)  // S_old * D
+        let dstBlockStride = dstShape[seqDim...].reduce(1, *)  // S_new * D
+
+        source.view(as: LogitsScalarType.self).withUnsafePointer { srcPtr, _, _ in
+            let dstView = destination.mutableView(as: LogitsScalarType.self)
+            dstView.withUnsafeMutablePointer { dstPtr, _, _ in
+                for block in 0..<numBlocks {
+                    let srcOff = block * srcBlockStride
+                    let dstOff = block * dstBlockStride
+                    dstPtr.advanced(by: dstOff).update(
+                        from: srcPtr.advanced(by: srcOff), count: copySize)
+                }
+            }
+        }
+    }
+
+    // MARK: - Helpers
+
+    private func zeroFill(_ array: inout NDArray) {
+        let count = array.shape.reduce(1, *)
+        let view = array.mutableView(as: LogitsScalarType.self)
+        // Inlined constant write — under -Onone, fillNDArray's
+        // `(Int) -> LogitsScalarType` closure is invoked per element (no inlining),
+        // which made zeroing the KV cache (~14.7M elements for a 32K-context
+        // Qwen3) take ~6 seconds per `reset()`. Direct loop keeps this in
+        // the few-ms range even unoptimized; under -O it lowers to memset.
+        view.withUnsafeMutablePointer { ptr, _, _ in
+            for i in 0..<count {
+                ptr[i] = 0
+            }
+        }
+    }
 }
 
 extension CoreAISequentialEngine {