Add native Swift CoreML runtime and assets

README.md

@@ -22,7 +22,7 @@ tags:
 
 This repository contains an early CoreML runtime bundle derived from [`datalab-to/surya-ocr-2`](https://huggingface.co/datalab-to/surya-ocr-2) at source commit `3b3d4cdf88d6928b0acdc75181b13206ea67c4a3`.
 
-It is intended for native Apple OCR experiments and future iOS/macOS demo-app work. It is **not yet a drop-in app runtime**: the model packages pass canary runtime gates, but host-side glue is still required.
+It is intended for native Apple OCR experiments and future iOS/macOS demo-app work. This revision includes a SwiftPM runtime package that runs the fixed-shape OCR canary path without Python.
 
 ## Included packages
 
@@ -33,7 +33,21 @@ It is intended for native Apple OCR experiments and future iOS/macOS demo-app wo
 | `surya_prefill_fp16_seq300_cache512.mlpackage` | Language prefill, logits, and initial cache | CoreML FP16 compute | fixed prefill length `300`, cache length `512` |
 | `surya_decode_step_fp16_cache512.mlpackage` | One-token cached decode step | CoreML FP16 compute | one token at a time, cache length `512` |
 
-Processor/tokenizer assets are included under `processor/`. Validation JSONs are under `validation/`.
+Processor/tokenizer assets are included under `processor/`. Native constants and embedding assets are under `native_assets/`. Validation JSONs are under `validation/`.
+
+## Native Swift package
+
+The SwiftPM package lives in `native/SuryaCoreMLRuntime`.
+
+```swift
+import SuryaCoreMLRuntime
+
+let runtime = try SuryaCoreMLRuntime(modelDirectory: modelDirectoryURL)
+let result = try runtime.generate(image: cgImage, maxNewTokens: 128, useInt8Vision: true)
+print(result.text ?? "")
+```
+
+The package currently handles fixed 512x512 Qwen image preprocessing, INT8/FP16 vision selection, fixed OCR prompt embeddings, image placeholder insertion, precomputed RoPE tables, generated-token embedding lookup, KV cache insertion, CoreML prefill/decode, and basic token decoding.
 
 ## Current validation
 
@@ -56,22 +70,20 @@ Total $42.00
 | CoreML prefill -> CoreML decode | 9/9 tokens match native; text `<p>Invoice ` |
 | CoreML FP16 vision -> CoreML prefill -> CoreML decode | 9/9 tokens match native; text `<p>Invoice ` |
 | CoreML INT8 vision -> CoreML prefill -> CoreML decode | 9/9 tokens match native; text `<p>Invoice ` |
+| Native Swift image -> CoreML INT8 vision -> CoreML prefill -> CoreML decode | 9/9 tokens match native; text `<p>Invoice ` |
 
 The INT8 vision package has mean absolute diff `0.021211756393313408` vs the PyTorch vision tower on the canary.
 
 ## What still lives in host code
 
-This is a runtime component bundle, not a complete Swift package yet. The current host/application layer must handle:
+The included Swift package is a fixed-shape native runtime for the current OCR path. The host app still owns:
 
-- tokenization
-- initial text token embedding lookup
-- image placeholder insertion
-- rotary position embedding generation
-- generated-token embedding lookup
-- full-attention KV cache insertion
-- stopping criteria and HTML/post-processing
+- providing a `CGImage` or already-preprocessed `pixel_values`
+- stopping criteria and max-token policy
+- UI and post-processing of generated OCR HTML
+- expanding beyond the current fixed 512x512 image / 300-token prefill shape
 
-The included `scripts/export_surya_coreml_runtime.py` shows the exact Python reference glue used for the validation runs.
+The included `scripts/export_surya_coreml_runtime.py` shows the Python reference glue used for validation. The included `native/SuryaCoreMLRuntime` package is the native implementation.
 
 ## Example: run the current validation harness
 
@@ -97,13 +109,31 @@ Expected canary output today:
 }
 ```
 
+## Example: run the native Swift smoke
+
+```bash
+cd native/SuryaCoreMLRuntime
+swift run surya-coreml-smoke \
+  --model-dir ../.. \
+  --image /path/to/512x512-document.png \
+  --max-tokens 8 \
+  --vision int8
+```
+
+Validated canary output:
+
+```text
+1039 2009 2046 2054 2047 2041 2035 2037 1979
+<p>Invoice 
+```
+
 ## Important limitations
 
 - Fixed-shape canary export: the current packages are specialized to the traced 512x512 sample preprocessing path, with `pixel_values [1024,1536]`, prefill length `300`, and full-attention cache length `512`.
 - The language prefill/decode packages are FP16 CoreML, not INT8/4-bit yet.
 - Only the vision tower has an INT8 package in this release.
 - This has not yet passed full `allenai/olmOCR-bench`.
-- This is not a complete native OCR app. It is the validated CoreML model core we will build the app around.
+- The Swift decoder is intentionally small and should be hardened against the full tokenizer before production use.
 
 ## Provenance
 

coreml_runtime_manifest.json

@@ -28,14 +28,12 @@
     "combined_prefill_decode_tokens_match": "9/9",
     "vision_fp16_prefill_decode_tokens_match": "9/9",
     "vision_int8_prefill_decode_tokens_match": "9/9",
+    "native_swift_vision_int8_prefill_decode_tokens_match": "9/9",
     "canary_text": "<p>Invoice "
   },
   "host_responsibilities": [
-    "tokenization",
-    "initial text token embedding lookup",
-    "image placeholder insertion",
-    "rotary position embedding generation",
-    "generated-token embedding lookup",
-    "full-attention KV cache insertion"
+    "provide a CGImage or preprocessed pixel_values",
+    "choose stopping criteria",
+    "post-process generated OCR HTML"
   ]
 }

native/SuryaCoreMLRuntime/Package.swift

@@ -0,0 +1,25 @@
+// swift-tools-version: 6.0
+import PackageDescription
+
+let package = Package(
+    name: "SuryaCoreMLRuntime",
+    platforms: [
+        .iOS(.v17),
+        .macOS(.v14),
+    ],
+    products: [
+        .library(name: "SuryaCoreMLRuntime", targets: ["SuryaCoreMLRuntime"]),
+        .executable(name: "surya-coreml-smoke", targets: ["SuryaCoreMLSmoke"]),
+    ],
+    targets: [
+        .target(name: "SuryaCoreMLRuntime"),
+        .executableTarget(
+            name: "SuryaCoreMLSmoke",
+            dependencies: ["SuryaCoreMLRuntime"]
+        ),
+        .testTarget(
+            name: "SuryaCoreMLRuntimeTests",
+            dependencies: ["SuryaCoreMLRuntime"]
+        ),
+    ]
+)

native/SuryaCoreMLRuntime/Sources/SuryaCoreMLRuntime/DecodeState.swift

@@ -0,0 +1,77 @@
+import CoreML
+import Foundation
+
+final class DecodeState {
+    var fullKeys: [MLMultiArray]
+    var fullValues: [MLMultiArray]
+    var convStates: [MLMultiArray]
+    var recurrentStates: [MLMultiArray]
+
+    init(prefillOutput: MLFeatureProvider, promptLength: Int) throws {
+        self.fullKeys = try (0..<SuryaConstants.fullAttentionLayers).map {
+            try prefillOutput.requiredArray("full_key_\($0)")
+        }
+        self.fullValues = try (0..<SuryaConstants.fullAttentionLayers).map {
+            try prefillOutput.requiredArray("full_value_\($0)")
+        }
+        self.convStates = try (0..<SuryaConstants.linearAttentionLayers).map {
+            try prefillOutput.requiredArray("conv_state_\($0)")
+        }
+        self.recurrentStates = try (0..<SuryaConstants.linearAttentionLayers).map {
+            try prefillOutput.requiredArray("recurrent_state_\($0)")
+        }
+
+        guard fullKeys.count == SuryaConstants.fullAttentionLayers,
+              fullValues.count == SuryaConstants.fullAttentionLayers,
+              convStates.count == SuryaConstants.linearAttentionLayers,
+              recurrentStates.count == SuryaConstants.linearAttentionLayers else {
+            throw SuryaRuntimeError.invalidShape("Missing cache state outputs")
+        }
+
+        _ = promptLength
+    }
+
+    func apply(decodeOutput: MLFeatureProvider, cacheIndex: Int) throws {
+        for layer in 0..<SuryaConstants.fullAttentionLayers {
+            let newKey = try decodeOutput.requiredArray("new_full_key_\(layer)")
+            let newValue = try decodeOutput.requiredArray("new_full_value_\(layer)")
+            try insertFullAttentionSlice(source: newKey, destination: fullKeys[layer], cacheIndex: cacheIndex)
+            try insertFullAttentionSlice(source: newValue, destination: fullValues[layer], cacheIndex: cacheIndex)
+        }
+
+        self.convStates = try (0..<SuryaConstants.linearAttentionLayers).map {
+            try decodeOutput.requiredArray("new_conv_state_\($0)")
+        }
+        self.recurrentStates = try (0..<SuryaConstants.linearAttentionLayers).map {
+            try decodeOutput.requiredArray("new_recurrent_state_\($0)")
+        }
+    }
+
+    private func insertFullAttentionSlice(source: MLMultiArray, destination: MLMultiArray, cacheIndex: Int) throws {
+        let dstShape = destination.shape.map(\.intValue)
+        guard dstShape.count == 4 else {
+            throw SuryaRuntimeError.invalidShape("Expected full cache rank 4, got \(dstShape)")
+        }
+        let batch = dstShape[0]
+        let heads = dstShape[1]
+        let cacheLength = dstShape[2]
+        let dim = dstShape[3]
+        guard batch == 1, cacheIndex >= 0, cacheIndex < cacheLength else {
+            throw SuryaRuntimeError.invalidShape("Bad cache index \(cacheIndex) for shape \(dstShape)")
+        }
+
+        let sourceCount = heads * dim
+        guard source.count == sourceCount || source.count == batch * sourceCount else {
+            throw SuryaRuntimeError.invalidShape("New KV slice count \(source.count), expected \(sourceCount)")
+        }
+
+        for head in 0..<heads {
+            for d in 0..<dim {
+                let srcIndex = head * dim + d
+                let dstIndex = ((head * cacheLength) + cacheIndex) * dim + d
+                let value = try source.readFloat(at: srcIndex)
+                try destination.writeFloat(value, at: dstIndex)
+            }
+        }
+    }
+}

native/SuryaCoreMLRuntime/Sources/SuryaCoreMLRuntime/MLMultiArrayUtilities.swift

@@ -0,0 +1,101 @@
+import CoreML
+import Foundation
+
+extension MLFeatureProvider {
+    func requiredArray(_ name: String) throws -> MLMultiArray {
+        guard let value = featureValue(for: name) else {
+            throw SuryaRuntimeError.missingFeature(name)
+        }
+        guard let array = value.multiArrayValue else {
+            throw SuryaRuntimeError.missingFeature("\(name) is not an MLMultiArray")
+        }
+        return array
+    }
+}
+
+extension MLMultiArray {
+    func floatPointer() -> UnsafeMutablePointer<Float> {
+        dataPointer.bindMemory(to: Float.self, capacity: count)
+    }
+
+    func halfPointer() -> UnsafeMutablePointer<Float16> {
+        dataPointer.bindMemory(to: Float16.self, capacity: count)
+    }
+
+    func readFloat(at linearIndex: Int) throws -> Float {
+        switch dataType {
+        case .float32:
+            return floatPointer()[linearIndex]
+        case .float16:
+            return Float(halfPointer()[linearIndex])
+        default:
+            throw SuryaRuntimeError.unsupportedDataType("Expected float16/float32, got \(dataType)")
+        }
+    }
+
+    func writeFloat(_ value: Float, at linearIndex: Int) throws {
+        switch dataType {
+        case .float32:
+            floatPointer()[linearIndex] = value
+        case .float16:
+            halfPointer()[linearIndex] = Float16(value)
+        default:
+            throw SuryaRuntimeError.unsupportedDataType("Expected float16/float32, got \(dataType)")
+        }
+    }
+
+    func argmaxLastDimension() throws -> Int32 {
+        guard count > 0 else {
+            throw SuryaRuntimeError.invalidShape("Cannot argmax empty MLMultiArray")
+        }
+        let last = shape.last?.intValue ?? count
+        guard last > 0, count % last == 0 else {
+            throw SuryaRuntimeError.invalidShape("Bad logits shape \(shape)")
+        }
+        let offset = count - last
+        var bestIndex = 0
+        var bestValue = try readFloat(at: offset)
+        if last > 1 {
+            for i in 1..<last {
+                let value = try readFloat(at: offset + i)
+                if value > bestValue {
+                    bestValue = value
+                    bestIndex = i
+                }
+            }
+        }
+        return Int32(bestIndex)
+    }
+}
+
+func product(_ shape: [Int]) -> Int {
+    shape.reduce(1, *)
+}
+
+func numberShape(_ shape: [Int]) -> [NSNumber] {
+    shape.map { NSNumber(value: $0) }
+}
+
+func contiguousFloat32Array(shape: [Int]) throws -> MLMultiArray {
+    try MLMultiArray(shape: numberShape(shape), dataType: .float32)
+}
+
+func loadRawFloat16AsFloat32(url: URL, shape: [Int]) throws -> MLMultiArray {
+    guard FileManager.default.fileExists(atPath: url.path) else {
+        throw SuryaRuntimeError.missingFile(url)
+    }
+    let data = try Data(contentsOf: url, options: [.mappedIfSafe])
+    let expectedScalars = product(shape)
+    guard data.count == expectedScalars * MemoryLayout<UInt16>.size else {
+        throw SuryaRuntimeError.invalidShape("\(url.lastPathComponent) has \(data.count) bytes, expected \(expectedScalars * 2)")
+    }
+    let array = try contiguousFloat32Array(shape: shape)
+    let dst = array.floatPointer()
+    data.withUnsafeBytes { raw in
+        let src = raw.bindMemory(to: UInt16.self)
+        for index in 0..<expectedScalars {
+            dst[index] = Float(Float16(bitPattern: src[index]))
+        }
+    }
+    return array
+}

native/SuryaCoreMLRuntime/Sources/SuryaCoreMLRuntime/QwenFixedImagePreprocessor.swift

@@ -0,0 +1,128 @@
+import CoreGraphics
+import CoreML
+import Foundation
+
+public struct QwenFixedImagePreprocessor: Sendable {
+    public let size: Int
+    public let patchSize: Int
+    public let temporalPatchSize: Int
+
+    public init(size: Int = 512, patchSize: Int = 16, temporalPatchSize: Int = 2) {
+        self.size = size
+        self.patchSize = patchSize
+        self.temporalPatchSize = temporalPatchSize
+    }
+
+    public func pixelValues(from image: CGImage) throws -> MLMultiArray {
+        let rgba = try resizedRGBA(image)
+        let patchesPerSide = size / patchSize
+        let mergeSize = 2
+        let mergedPatchesPerSide = patchesPerSide / mergeSize
+        let patchCount = patchesPerSide * patchesPerSide
+        let patchWidth = 3 * temporalPatchSize * patchSize * patchSize
+        let output = try MLMultiArray(shape: [NSNumber(value: patchCount), NSNumber(value: patchWidth)], dataType: .float32)
+        let dst = output.floatPointer()
+
+        var patchIndex = 0
+        for blockY in 0..<mergedPatchesPerSide {
+            for blockX in 0..<mergedPatchesPerSide {
+                for mergeY in 0..<mergeSize {
+                    for mergeX in 0..<mergeSize {
+                        let patchY = blockY * mergeSize + mergeY
+                        let patchX = blockX * mergeSize + mergeX
+                        let patchOffset = patchIndex * patchWidth
+                        var cursor = patchOffset
+                        for channel in 0..<3 {
+                            for _ in 0..<temporalPatchSize {
+                                for y in 0..<patchSize {
+                                    let imageY = patchY * patchSize + y
+                                    for x in 0..<patchSize {
+                                        let imageX = patchX * patchSize + x
+                                        let pixelOffset = (imageY * size + imageX) * 4
+                                        let raw: UInt8
+                                        switch channel {
+                                        case 0: raw = rgba[pixelOffset]
+                                        case 1: raw = rgba[pixelOffset + 1]
+                                        default: raw = rgba[pixelOffset + 2]
+                                        }
+                                        dst[cursor] = (Float(raw) / 255.0 - 0.5) / 0.5
+                                        cursor += 1
+                                    }
+                                }
+                            }
+                        }
+                        patchIndex += 1
+                    }
+                }
+            }
+        }
+
+        return output
+    }
+
+    private func resizedRGBA(_ image: CGImage) throws -> [UInt8] {
+        if image.width == size,
+           image.height == size,
+           image.bitsPerComponent == 8,
+           image.bitsPerPixel == 32,
+           let dataProvider = image.dataProvider,
+           let data = dataProvider.data,
+           let source = CFDataGetBytePtr(data) {
+            let length = CFDataGetLength(data)
+            let bytesPerRow = image.bytesPerRow
+            if length >= bytesPerRow * size {
+                var pixels = [UInt8](repeating: 255, count: size * size * 4)
+                for y in 0..<size {
+                    let srcRow = source + y * bytesPerRow
+                    let dstRow = y * size * 4
+                    for x in 0..<size {
+                        let srcOffset = x * 4
+                        let dstOffset = dstRow + x * 4
+                        pixels[dstOffset] = srcRow[srcOffset]
+                        pixels[dstOffset + 1] = srcRow[srcOffset + 1]
+                        pixels[dstOffset + 2] = srcRow[srcOffset + 2]
+                        pixels[dstOffset + 3] = 255
+                    }
+                }
+                return pixels
+            }
+        }
+
+        var pixels = [UInt8](repeating: 255, count: size * size * 4)
+        let bitmapInfo = CGImageAlphaInfo.noneSkipLast.rawValue | CGBitmapInfo.byteOrder32Big.rawValue
+        guard let colorSpace = CGColorSpace(name: CGColorSpace.linearSRGB) ?? CGColorSpace(name: CGColorSpace.sRGB),
+              let context = CGContext(
+                  data: &pixels,
+                  width: size,
+                  height: size,
+                  bitsPerComponent: 8,
+                  bytesPerRow: size * 4,
+                  space: colorSpace,
+                  bitmapInfo: bitmapInfo
+              ) else {
+            throw SuryaRuntimeError.invalidShape("Could not create CGContext for image preprocessing")
+        }
+
+        context.interpolationQuality = .none
+        context.setFillColor(CGColor(red: 1, green: 1, blue: 1, alpha: 1))
+        context.fill(CGRect(x: 0, y: 0, width: size, height: size))
+
+        let sourceWidth = CGFloat(image.width)
+        let sourceHeight = CGFloat(image.height)
+        let scale = min(CGFloat(size) / sourceWidth, CGFloat(size) / sourceHeight)
+        let drawWidth = sourceWidth * scale
+        let drawHeight = sourceHeight * scale
+        let drawRect = CGRect(
+            x: (CGFloat(size) - drawWidth) / 2,
+            y: (CGFloat(size) - drawHeight) / 2,
+            width: drawWidth,
+            height: drawHeight
+        )
+        context.saveGState()
+        context.translateBy(x: 0, y: CGFloat(size))
+        context.scaleBy(x: 1, y: -1)
+        context.draw(image, in: drawRect)
+        context.restoreGState()
+        return pixels
+    }
+}

native/SuryaCoreMLRuntime/Sources/SuryaCoreMLRuntime/SuryaCoreMLRuntime.swift

@@ -0,0 +1,152 @@
+import CoreGraphics
+import CoreML
+import Foundation
+
+public enum SuryaRuntimeError: Error, CustomStringConvertible {
+    case missingFile(URL)
+    case malformedConstants(String)
+    case missingFeature(String)
+    case unsupportedDataType(String)
+    case invalidShape(String)
+    case generationLimit(String)
+
+    public var description: String {
+        switch self {
+        case .missingFile(let url): return "Missing file: \(url.path)"
+        case .malformedConstants(let message): return "Malformed constants: \(message)"
+        case .missingFeature(let name): return "Missing CoreML feature: \(name)"
+        case .unsupportedDataType(let message): return "Unsupported data type: \(message)"
+        case .invalidShape(let message): return "Invalid shape: \(message)"
+        case .generationLimit(let message): return "Generation limit: \(message)"
+        }
+    }
+}
+
+public struct SuryaGenerationResult: Sendable {
+    public let tokenIDs: [Int32]
+    public let text: String?
+
+    public init(tokenIDs: [Int32], text: String? = nil) {
+        self.tokenIDs = tokenIDs
+        self.text = text
+    }
+}
+
+public final class SuryaCoreMLRuntime {
+    private let bundle: SuryaModelBundle
+    private let constants: SuryaNativeConstants
+    private let arrays: SuryaNativeArrays
+    private let preprocessor: QwenFixedImagePreprocessor
+    private let tokenDecoder: TokenDecoder
+
+    public init(modelDirectory: URL, computeUnits: MLComputeUnits = .cpuOnly) throws {
+        self.bundle = try SuryaModelBundle(modelDirectory: modelDirectory, computeUnits: computeUnits)
+        self.constants = try SuryaNativeConstants.load(from: modelDirectory.appendingPathComponent("native_assets/surya_native_constants.json"))
+        self.arrays = try SuryaNativeArrays(modelDirectory: modelDirectory, constants: constants)
+        self.preprocessor = QwenFixedImagePreprocessor()
+        self.tokenDecoder = try TokenDecoder(modelDirectory: modelDirectory)
+    }
+
+    public func generate(
+        image: CGImage,
+        maxNewTokens: Int = 128,
+        useInt8Vision: Bool = true,
+        stopTokenIDs: Set<Int32> = []
+    ) throws -> SuryaGenerationResult {
+        let pixelValues = try preprocessor.pixelValues(from: image)
+        return try generate(pixelValues: pixelValues, maxNewTokens: maxNewTokens, useInt8Vision: useInt8Vision, stopTokenIDs: stopTokenIDs)
+    }
+
+    public func generate(
+        pixelValues: MLMultiArray,
+        maxNewTokens: Int = 128,
+        useInt8Vision: Bool = true,
+        stopTokenIDs: Set<Int32> = []
+    ) throws -> SuryaGenerationResult {
+        let vision = useInt8Vision ? bundle.visionInt8 : bundle.visionFP16
+        let visionOutput = try vision.prediction(from: ["pixel_values": MLFeatureValue(multiArray: pixelValues)])
+        let imageEmbeds = try visionOutput.requiredArray("image_embeds")
+
+        let prefillInputs = try buildPrefillInputs(imageEmbeds: imageEmbeds)
+        let prefillOutput = try bundle.prefill.prediction(from: prefillInputs)
+        var generated: [Int32] = [try argmax(featureProvider: prefillOutput, name: "logits")]
+
+        let state = try DecodeState(prefillOutput: prefillOutput, promptLength: constants.promptLength)
+        var cacheLength = constants.promptLength
+        var currentToken = generated[0]
+
+        for _ in 0..<maxNewTokens {
+            if stopTokenIDs.contains(currentToken) {
+                break
+            }
+            if cacheLength >= constants.maxCacheLength {
+                throw SuryaRuntimeError.generationLimit("cacheLength \(cacheLength) reached max \(constants.maxCacheLength)")
+            }
+
+            let decodeInputs = try buildDecodeInputs(currentToken: currentToken, cacheLength: cacheLength, state: state)
+            let decodeOutput = try bundle.decodeStep.prediction(from: decodeInputs)
+            let nextToken = try argmax(featureProvider: decodeOutput, name: "logits")
+
+            try state.apply(decodeOutput: decodeOutput, cacheIndex: cacheLength)
+            cacheLength += 1
+            currentToken = nextToken
+            generated.append(nextToken)
+        }
+
+        return SuryaGenerationResult(tokenIDs: generated, text: tokenDecoder.decode(generated))
+    }
+
+    private func buildPrefillInputs(imageEmbeds: MLMultiArray) throws -> [String: MLFeatureValue] {
+        let inputsEmbeds = try arrays.makePrefillInputsEmbeds(imageEmbeds: imageEmbeds, imageTokenIndices: constants.imageTokenIndices)
+        return [
+            "inputs_embeds": MLFeatureValue(multiArray: inputsEmbeds),
+            "cos": MLFeatureValue(multiArray: arrays.prefillCos),
+            "sin": MLFeatureValue(multiArray: arrays.prefillSin),
+        ]
+    }
+
+    private func buildDecodeInputs(currentToken: Int32, cacheLength: Int, state: DecodeState) throws -> [String: MLFeatureValue] {
+        var features: [String: MLFeatureValue] = [
+            "inputs_embeds": MLFeatureValue(multiArray: try arrays.embedding(for: currentToken)),
+            "cos": MLFeatureValue(multiArray: try arrays.decodeCos(position: cacheLength)),
+            "sin": MLFeatureValue(multiArray: try arrays.decodeSin(position: cacheLength)),
+            "attention_mask": MLFeatureValue(multiArray: try attentionMask(cacheLength: cacheLength)),
+        ]
+
+        for index in 0..<SuryaConstants.fullAttentionLayers {
+            features["full_key_\(index)"] = MLFeatureValue(multiArray: state.fullKeys[index])
+            features["full_value_\(index)"] = MLFeatureValue(multiArray: state.fullValues[index])
+        }
+        for index in 0..<SuryaConstants.linearAttentionLayers {
+            features["conv_state_\(index)"] = MLFeatureValue(multiArray: state.convStates[index])
+            features["recurrent_state_\(index)"] = MLFeatureValue(multiArray: state.recurrentStates[index])
+        }
+        return features
+    }
+
+    private func attentionMask(cacheLength: Int) throws -> MLMultiArray {
+        let width = constants.maxCacheLength + 1
+        let array = try MLMultiArray(shape: [1, 1, 1, NSNumber(value: width)], dataType: .float32)
+        let values = array.floatPointer()
+        let negative = -Float.greatestFiniteMagnitude
+        for index in 0..<width {
+            values[index] = negative
+        }
+        for index in 0..<cacheLength {
+            values[index] = 0
+        }
+        values[width - 1] = 0
+        return array
+    }
+
+    private func argmax(featureProvider: MLFeatureProvider, name: String) throws -> Int32 {
+        let logits = try featureProvider.requiredArray(name)
+        return try logits.argmaxLastDimension()
+    }
+}
+
+enum SuryaConstants {
+    static let hiddenSize = 1024
+    static let fullAttentionLayers = 6
+    static let linearAttentionLayers = 18
+}

native/SuryaCoreMLRuntime/Sources/SuryaCoreMLRuntime/SuryaModelBundle.swift

@@ -0,0 +1,35 @@
+import CoreML
+import Foundation
+
+final class SuryaModelBundle {
+    let visionFP16: MLModel
+    let visionInt8: MLModel
+    let prefill: MLModel
+    let decodeStep: MLModel
+
+    init(modelDirectory: URL, computeUnits: MLComputeUnits) throws {
+        let configuration = MLModelConfiguration()
+        configuration.computeUnits = computeUnits
+
+        self.visionFP16 = try Self.loadPackage("surya_vision_fp16.mlpackage", in: modelDirectory, configuration: configuration)
+        self.visionInt8 = try Self.loadPackage("surya_vision_int8.mlpackage", in: modelDirectory, configuration: configuration)
+        self.prefill = try Self.loadPackage("surya_prefill_fp16_seq300_cache512.mlpackage", in: modelDirectory, configuration: configuration)
+        self.decodeStep = try Self.loadPackage("surya_decode_step_fp16_cache512.mlpackage", in: modelDirectory, configuration: configuration)
+    }
+
+    private static func loadPackage(_ name: String, in directory: URL, configuration: MLModelConfiguration) throws -> MLModel {
+        let url = directory.appendingPathComponent(name)
+        guard FileManager.default.fileExists(atPath: url.path) else {
+            throw SuryaRuntimeError.missingFile(url)
+        }
+        let compiledURL = try MLModel.compileModel(at: url)
+        return try MLModel(contentsOf: compiledURL, configuration: configuration)
+    }
+}
+
+extension MLModel {
+    func prediction(from features: [String: MLFeatureValue]) throws -> MLFeatureProvider {
+        let provider = try MLDictionaryFeatureProvider(dictionary: features)
+        return try prediction(from: provider)
+    }
+}

native/SuryaCoreMLRuntime/Sources/SuryaCoreMLRuntime/SuryaNativeArrays.swift

@@ -0,0 +1,109 @@
+import CoreML
+import Foundation
+
+final class SuryaNativeArrays {
+    private let constants: SuryaNativeConstants
+    let prefillTextEmbedsBase: MLMultiArray
+    let prefillCos: MLMultiArray
+    let prefillSin: MLMultiArray
+    private let decodeCosTable: MLMultiArray
+    private let decodeSinTable: MLMultiArray
+    private let tokenEmbedding: MLMultiArray
+
+    init(modelDirectory: URL, constants: SuryaNativeConstants) throws {
+        self.constants = constants
+        let assets = modelDirectory.appendingPathComponent("native_assets")
+        self.prefillTextEmbedsBase = try loadRawFloat16AsFloat32(
+            url: assets.appendingPathComponent(constants.files.prefillTextEmbedsBase),
+            shape: constants.shapes.prefillTextEmbedsBase
+        )
+        self.prefillCos = try loadRawFloat16AsFloat32(
+            url: assets.appendingPathComponent(constants.files.prefillCos),
+            shape: constants.shapes.prefillCos
+        )
+        self.prefillSin = try loadRawFloat16AsFloat32(
+            url: assets.appendingPathComponent(constants.files.prefillSin),
+            shape: constants.shapes.prefillSin
+        )
+        self.decodeCosTable = try loadRawFloat16AsFloat32(
+            url: assets.appendingPathComponent(constants.files.decodeCos),
+            shape: constants.shapes.decodeCos
+        )
+        self.decodeSinTable = try loadRawFloat16AsFloat32(
+            url: assets.appendingPathComponent(constants.files.decodeSin),
+            shape: constants.shapes.decodeSin
+        )
+        self.tokenEmbedding = try loadRawFloat16AsFloat32(
+            url: assets.appendingPathComponent(constants.files.tokenEmbedding),
+            shape: constants.shapes.tokenEmbedding
+        )
+    }
+
+    func makePrefillInputsEmbeds(imageEmbeds: MLMultiArray, imageTokenIndices: [Int]) throws -> MLMultiArray {
+        let shape = constants.shapes.prefillTextEmbedsBase
+        guard shape == [1, constants.promptLength, constants.hiddenSize] else {
+            throw SuryaRuntimeError.invalidShape("Unexpected prefill base shape \(shape)")
+        }
+        guard imageTokenIndices.count == constants.shapes.imageEmbeds[0] else {
+            throw SuryaRuntimeError.invalidShape("Image token count \(imageTokenIndices.count) does not match image embeds shape \(constants.shapes.imageEmbeds)")
+        }
+
+        let result = try contiguousFloat32Array(shape: shape)
+        let srcBase = prefillTextEmbedsBase.floatPointer()
+        let dst = result.floatPointer()
+        for index in 0..<prefillTextEmbedsBase.count {
+            dst[index] = srcBase[index]
+        }
+
+        let hidden = constants.hiddenSize
+        for imageIndex in 0..<imageTokenIndices.count {
+            let tokenIndex = imageTokenIndices[imageIndex]
+            let dstOffset = tokenIndex * hidden
+            let srcOffset = imageIndex * hidden
+            for hiddenIndex in 0..<hidden {
+                dst[dstOffset + hiddenIndex] = try imageEmbeds.readFloat(at: srcOffset + hiddenIndex)
+            }
+        }
+        return result
+    }
+
+    func embedding(for tokenID: Int32) throws -> MLMultiArray {
+        let token = Int(tokenID)
+        guard token >= 0, token < constants.vocabSize else {
+            throw SuryaRuntimeError.invalidShape("Token id \(tokenID) outside vocab \(constants.vocabSize)")
+        }
+        let hidden = constants.hiddenSize
+        let result = try contiguousFloat32Array(shape: [1, 1, hidden])
+        let src = tokenEmbedding.floatPointer()
+        let dst = result.floatPointer()
+        let offset = token * hidden
+        for i in 0..<hidden {
+            dst[i] = src[offset + i]
+        }
+        return result
+    }
+
+    func decodeCos(position: Int) throws -> MLMultiArray {
+        try decodeTableSlice(decodeCosTable, position: position)
+    }
+
+    func decodeSin(position: Int) throws -> MLMultiArray {
+        try decodeTableSlice(decodeSinTable, position: position)
+    }
+
+    private func decodeTableSlice(_ table: MLMultiArray, position: Int) throws -> MLMultiArray {
+        let shape = constants.shapes.decodeCos
+        guard shape.count == 3, position >= 0, position < shape[1] else {
+            throw SuryaRuntimeError.invalidShape("Bad decode position \(position) for shape \(shape)")
+        }
+        let width = shape[2]
+        let result = try contiguousFloat32Array(shape: [1, 1, width])
+        let src = table.floatPointer()
+        let dst = result.floatPointer()
+        let offset = position * width
+        for i in 0..<width {
+            dst[i] = src[offset + i]
+        }
+        return result
+    }
+}

native/SuryaCoreMLRuntime/Sources/SuryaCoreMLRuntime/SuryaNativeConstants.swift

@@ -0,0 +1,83 @@
+import Foundation
+
+struct SuryaNativeConstants: Decodable {
+    struct ShapeSet: Decodable {
+        let prefillTextEmbedsBase: [Int]
+        let prefillCos: [Int]
+        let prefillSin: [Int]
+        let decodeCos: [Int]
+        let decodeSin: [Int]
+        let tokenEmbedding: [Int]
+        let pixelValues: [Int]
+        let imageEmbeds: [Int]
+
+        enum CodingKeys: String, CodingKey {
+            case prefillTextEmbedsBase = "prefill_text_embeds_base"
+            case prefillCos = "prefill_cos"
+            case prefillSin = "prefill_sin"
+            case decodeCos = "decode_cos"
+            case decodeSin = "decode_sin"
+            case tokenEmbedding = "token_embedding"
+            case pixelValues = "pixel_values"
+            case imageEmbeds = "image_embeds"
+        }
+    }
+
+    struct FileSet: Decodable {
+        let prefillTextEmbedsBase: String
+        let prefillCos: String
+        let prefillSin: String
+        let decodeCos: String
+        let decodeSin: String
+        let tokenEmbedding: String
+
+        enum CodingKeys: String, CodingKey {
+            case prefillTextEmbedsBase = "prefill_text_embeds_base"
+            case prefillCos = "prefill_cos"
+            case prefillSin = "prefill_sin"
+            case decodeCos = "decode_cos"
+            case decodeSin = "decode_sin"
+            case tokenEmbedding = "token_embedding"
+        }
+    }
+
+    let modelID: String
+    let prompt: String
+    let inputIDs: [Int32]
+    let attentionMask: [Int]
+    let mmTokenTypeIDs: [Int]
+    let imageGridTHW: [Int]
+    let imageTokenIndices: [Int]
+    let shapes: ShapeSet
+    let files: FileSet
+
+    enum CodingKeys: String, CodingKey {
+        case modelID = "model_id"
+        case prompt
+        case inputIDs = "input_ids"
+        case attentionMask = "attention_mask"
+        case mmTokenTypeIDs = "mm_token_type_ids"
+        case imageGridTHW = "image_grid_thw"
+        case imageTokenIndices = "image_token_indices"
+        case shapes
+        case files
+    }
+
+    var promptLength: Int { inputIDs.count }
+    var maxCacheLength: Int { shapes.decodeCos.count >= 2 ? shapes.decodeCos[1] : 512 }
+    var vocabSize: Int { shapes.tokenEmbedding[0] }
+    var hiddenSize: Int { shapes.tokenEmbedding[1] }
+
+    static func load(from url: URL) throws -> SuryaNativeConstants {
+        guard FileManager.default.fileExists(atPath: url.path) else {
+            throw SuryaRuntimeError.missingFile(url)
+        }
+        let data = try Data(contentsOf: url)
+        let decoder = JSONDecoder()
+        do {
+            return try decoder.decode(SuryaNativeConstants.self, from: data)
+        } catch {
+            throw SuryaRuntimeError.malformedConstants(error.localizedDescription)
+        }
+    }
+}

native/SuryaCoreMLRuntime/Sources/SuryaCoreMLRuntime/TokenDecoder.swift

@@ -0,0 +1,65 @@
+import Foundation
+
+final class TokenDecoder {
+    private let idToToken: [String]
+
+    init(modelDirectory: URL) throws {
+        let url = modelDirectory.appendingPathComponent("native_assets/id_to_token.json")
+        guard FileManager.default.fileExists(atPath: url.path) else {
+            throw SuryaRuntimeError.missingFile(url)
+        }
+        let data = try Data(contentsOf: url, options: [.mappedIfSafe])
+        self.idToToken = try JSONDecoder().decode([String].self, from: data)
+    }
+
+    func decode(_ tokenIDs: [Int32]) -> String {
+        var output = ""
+        var byteBuffer: [UInt8] = []
+
+        func flushBytes() {
+            guard !byteBuffer.isEmpty else { return }
+            output += String(decoding: byteBuffer, as: UTF8.self)
+            byteBuffer.removeAll(keepingCapacity: true)
+        }
+
+        for id in tokenIDs {
+            let index = Int(id)
+            guard index >= 0, index < idToToken.count else { continue }
+            let token = idToToken[index]
+            if token.isEmpty { continue }
+
+            if let byte = Self.byteFallback(token) {
+                byteBuffer.append(byte)
+                continue
+            }
+
+            flushBytes()
+            output += Self.cleanToken(token)
+        }
+
+        flushBytes()
+        return output
+    }
+
+    private static func cleanToken(_ token: String) -> String {
+        if token == "Ġ" || token == "▁" {
+            return " "
+        }
+        var value = token
+        if value.hasPrefix("Ġ") || value.hasPrefix("▁") {
+            value.removeFirst()
+            value = " " + value
+        }
+        value = value.replacingOccurrences(of: "Ċ", with: "\n")
+        return value
+    }
+
+    private static func byteFallback(_ token: String) -> UInt8? {
+        guard token.hasPrefix("<0x"), token.hasSuffix(">") else {
+            return nil
+        }
+        let start = token.index(token.startIndex, offsetBy: 3)
+        let end = token.index(before: token.endIndex)
+        return UInt8(token[start..<end], radix: 16)
+    }
+}

native/SuryaCoreMLRuntime/Sources/SuryaCoreMLSmoke/main.swift

@@ -0,0 +1,130 @@
+import CoreGraphics
+import CoreML
+import Foundation
+import ImageIO
+import SuryaCoreMLRuntime
+
+struct Arguments {
+    let modelDirectory: URL
+    let imageURL: URL?
+    let pixelValuesURL: URL?
+    let dumpPixelValuesURL: URL?
+    let maxTokens: Int
+    let useInt8Vision: Bool
+}
+
+func parseArguments() throws -> Arguments {
+    var modelDirectory: URL?
+    var imageURL: URL?
+    var pixelValuesURL: URL?
+    var dumpPixelValuesURL: URL?
+    var maxTokens = 32
+    var useInt8Vision = true
+
+    var iterator = CommandLine.arguments.dropFirst().makeIterator()
+    while let arg = iterator.next() {
+        switch arg {
+        case "--model-dir":
+            guard let value = iterator.next() else { throw SuryaRuntimeError.malformedConstants("Missing --model-dir value") }
+            modelDirectory = URL(fileURLWithPath: value)
+        case "--image":
+            guard let value = iterator.next() else { throw SuryaRuntimeError.malformedConstants("Missing --image value") }
+            imageURL = URL(fileURLWithPath: value)
+        case "--pixel-values-fp16":
+            guard let value = iterator.next() else { throw SuryaRuntimeError.malformedConstants("Missing --pixel-values-fp16 value") }
+            pixelValuesURL = URL(fileURLWithPath: value)
+        case "--dump-pixel-values-fp16":
+            guard let value = iterator.next() else { throw SuryaRuntimeError.malformedConstants("Missing --dump-pixel-values-fp16 value") }
+            dumpPixelValuesURL = URL(fileURLWithPath: value)
+        case "--max-tokens":
+            guard let value = iterator.next(), let parsed = Int(value) else {
+                throw SuryaRuntimeError.malformedConstants("Bad --max-tokens value")
+            }
+            maxTokens = parsed
+        case "--vision":
+            guard let value = iterator.next() else { throw SuryaRuntimeError.malformedConstants("Missing --vision value") }
+            useInt8Vision = value.lowercased() != "fp16"
+        default:
+            throw SuryaRuntimeError.malformedConstants("Unknown argument \(arg)")
+        }
+    }
+
+    guard let modelDirectory else {
+        throw SuryaRuntimeError.malformedConstants("Usage: surya-coreml-smoke --model-dir /path/to/repo --image /path/to/image [--max-tokens 32] [--vision int8|fp16]")
+    }
+    if imageURL == nil && pixelValuesURL == nil {
+        throw SuryaRuntimeError.malformedConstants("Pass either --image or --pixel-values-fp16")
+    }
+
+    return Arguments(
+        modelDirectory: modelDirectory,
+        imageURL: imageURL,
+        pixelValuesURL: pixelValuesURL,
+        dumpPixelValuesURL: dumpPixelValuesURL,
+        maxTokens: maxTokens,
+        useInt8Vision: useInt8Vision
+    )
+}
+
+func loadImage(_ url: URL) throws -> CGImage {
+    guard let source = CGImageSourceCreateWithURL(url as CFURL, nil),
+          let image = CGImageSourceCreateImageAtIndex(source, 0, nil) else {
+        throw SuryaRuntimeError.missingFile(url)
+    }
+    return image
+}
+
+func loadPixelValues(_ url: URL) throws -> MLMultiArray {
+    let data = try Data(contentsOf: url, options: [.mappedIfSafe])
+    let shape = [1024, 1536]
+    let expected = shape.reduce(1, *) * MemoryLayout<UInt16>.size
+    guard data.count == expected else {
+        throw SuryaRuntimeError.invalidShape("Expected \(expected) bytes, got \(data.count)")
+    }
+    let array = try MLMultiArray(shape: shape.map(NSNumber.init(value:)), dataType: .float32)
+    let dst = array.dataPointer.bindMemory(to: Float.self, capacity: shape.reduce(1, *))
+    data.withUnsafeBytes { raw in
+        let src = raw.bindMemory(to: UInt16.self)
+        for i in 0..<src.count {
+            dst[i] = Float(Float16(bitPattern: src[i]))
+        }
+    }
+    return array
+}
+
+func dumpPixelValues(_ array: MLMultiArray, to url: URL) throws {
+    var data = Data(capacity: array.count * MemoryLayout<UInt16>.size)
+    for index in 0..<array.count {
+        let value = array.dataPointer.bindMemory(to: Float.self, capacity: array.count)[index]
+        var bits = Float16(value).bitPattern
+        withUnsafeBytes(of: &bits) { data.append(contentsOf: $0) }
+    }
+    try data.write(to: url)
+}
+
+do {
+    let args = try parseArguments()
+    let runtime = try SuryaCoreMLRuntime(modelDirectory: args.modelDirectory)
+    let result: SuryaGenerationResult
+    if let pixelValuesURL = args.pixelValuesURL {
+        let pixelValues = try loadPixelValues(pixelValuesURL)
+        result = try runtime.generate(pixelValues: pixelValues, maxNewTokens: args.maxTokens, useInt8Vision: args.useInt8Vision)
+    } else if let imageURL = args.imageURL {
+        let image = try loadImage(imageURL)
+        let preprocessor = QwenFixedImagePreprocessor()
+        let pixelValues = try preprocessor.pixelValues(from: image)
+        if let dumpURL = args.dumpPixelValuesURL {
+            try dumpPixelValues(pixelValues, to: dumpURL)
+        }
+        result = try runtime.generate(pixelValues: pixelValues, maxNewTokens: args.maxTokens, useInt8Vision: args.useInt8Vision)
+    } else {
+        throw SuryaRuntimeError.malformedConstants("Pass either --image or --pixel-values-fp16")
+    }
+    print(result.tokenIDs.map(String.init).joined(separator: " "))
+    if let text = result.text {
+        print(text)
+    }
+} catch {
+    fputs("surya-coreml-smoke failed: \(error)\n", stderr)
+    exit(1)
+}

native/SuryaCoreMLRuntime/Tests/SuryaCoreMLRuntimeTests/SuryaCoreMLRuntimeTests.swift

@@ -0,0 +1,28 @@
+import CoreML
+import Testing
+@testable import SuryaCoreMLRuntime
+
+@Test func qwenFixedImagePreprocessorShape() throws {
+    let width = 32
+    let height = 32
+    var pixels = [UInt8](repeating: 255, count: width * height * 4)
+    for i in stride(from: 0, to: pixels.count, by: 4) {
+        pixels[i] = 0
+        pixels[i + 1] = 0
+        pixels[i + 2] = 0
+        pixels[i + 3] = 255
+    }
+    let colorSpace = CGColorSpace(name: CGColorSpace.sRGB)!
+    let context = CGContext(
+        data: &pixels,
+        width: width,
+        height: height,
+        bitsPerComponent: 8,
+        bytesPerRow: width * 4,
+        space: colorSpace,
+        bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue
+    )!
+    let image = context.makeImage()!
+    let values = try QwenFixedImagePreprocessor().pixelValues(from: image)
+    #expect(values.shape.map(\.intValue) == [1024, 1536])
+}

native_assets/canary_pixel_values_fp16.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe20bcb89c45b610cc3912c9f2b83e5bb7e374370fefafbb86ab185ef318fdbe
+size 3145728

native_assets/decode_cos_fp16.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9dfc5caf59d669a7acc909daae1a250a260802b5f123bbedc3038010cdb459c7
+size 65536

native_assets/decode_sin_fp16.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:67f1a0febb226f6521ed34c32c5669c0b8356bcfedd5bf27d27b006ebd63eeab
+size 65536

native_assets/prefill_cos_fp16.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a44a96a45de3e01c5036dc145460e81a49a40cc101b8fccca93255a5f401bfde
+size 38400

native_assets/prefill_sin_fp16.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:367b5007b74b28b90677d49dcc8813a49d67a2a3e5be8682da42facb1776496c
+size 38400

native_assets/prefill_text_embeds_base_fp16.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ee2b84e91c671347cc07c2f614821744728b6e2c3938d0cc7f618a41f104147
+size 614400

native_assets/surya_native_constants.json

@@ -0,0 +1,1222 @@
+{
+  "model_id": "datalab-to/surya-ocr-2",
+  "prompt": "OCR this image to HTML.",
+  "input_ids": [
+    1,
+    2053,
+    2051,
+    2037,
+    2050,
+    1957,
+    9,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    11,
+    10,
+    2015,
+    2003,
+    2018,
+    1979,
+    2052,
+    2040,
+    2041,
+    2051,
+    1979,
+    2041,
+    2045,
+    2033,
+    2039,
+    2037,
+    1979,
+    2052,
+    2047,
+    1979,
+    2008,
+    2020,
+    2013,
+    2012,
+    1992,
+    2,
+    1957,
+    1,
+    2033,
+    2051,
+    2051,
+    2041,
+    2051,
+    2052,
+    2033,
+    2046,
+    2052,
+    1957
+  ],
+  "attention_mask": [
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1
+  ],
+  "mm_token_type_ids": [
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0
+  ],
+  "image_grid_thw": [
+    1,
+    32,
+    32
+  ],
+  "image_token_indices": [
+    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,
+    194,
+    195,
+    196,
+    197,
+    198,
+    199,
+    200,
+    201,
+    202,
+    203,
+    204,
+    205,
+    206,
+    207,
+    208,
+    209,
+    210,
+    211,
+    212,
+    213,
+    214,
+    215,
+    216,
+    217,
+    218,
+    219,
+    220,
+    221,
+    222,
+    223,
+    224,
+    225,
+    226,
+    227,
+    228,
+    229,
+    230,
+    231,
+    232,
+    233,
+    234,
+    235,
+    236,
+    237,
+    238,
+    239,
+    240,
+    241,
+    242,
+    243,
+    244,
+    245,
+    246,
+    247,
+    248,
+    249,
+    250,
+    251,
+    252,
+    253,
+    254,
+    255,
+    256,
+    257,
+    258,
+    259,
+    260,
+    261,
+    262
+  ],
+  "shapes": {
+    "prefill_text_embeds_base": [
+      1,
+      300,
+      1024
+    ],
+    "prefill_cos": [
+      1,
+      300,
+      64
+    ],
+    "prefill_sin": [
+      1,
+      300,
+      64
+    ],
+    "decode_cos": [
+      1,
+      512,
+      64
+    ],
+    "decode_sin": [
+      1,
+      512,
+      64
+    ],
+    "token_embedding": [
+      65425,
+      1024
+    ],
+    "pixel_values": [
+      1024,
+      1536
+    ],
+    "image_embeds": [
+      256,
+      1024
+    ]
+  },
+  "dtype": "float16 little-endian raw binaries",
+  "files": {
+    "prefill_text_embeds_base": "prefill_text_embeds_base_fp16.bin",
+    "prefill_cos": "prefill_cos_fp16.bin",
+    "prefill_sin": "prefill_sin_fp16.bin",
+    "decode_cos": "decode_cos_fp16.bin",
+    "decode_sin": "decode_sin_fp16.bin",
+    "token_embedding": "token_embedding_fp16.bin",
+    "canary_pixel_values": "canary_pixel_values_fp16.bin"
+  }
+}

native_assets/swift_pixel_values_fp16.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:42632cf4f35becf36e42c03f6ad4a78f08e7cdeca8635f9283fe95be82d49b79
+size 3145728

native_assets/token_embedding_fp16.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a00af278f8f6d886c32a9397c7f3a4fea2de987b105131d835d6461131829aba
+size 133990400

scripts/materialize_surya_native_assets.py

@@ -0,0 +1,129 @@
+#!/usr/bin/env python3
+import argparse
+import importlib.util
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from transformers import AutoProcessor
+
+
+def load_runtime_module(path: Path):
+    spec = importlib.util.spec_from_file_location("surya_coreml_runtime_export", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def write_fp16(path: Path, tensor: torch.Tensor) -> None:
+    array = tensor.detach().cpu().to(torch.float16).numpy()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_bytes(array.tobytes(order="C"))
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-id", default="datalab-to/surya-ocr-2")
+    parser.add_argument("--runtime-script", type=Path, default=Path(".context/scripts/export_surya_coreml_runtime.py"))
+    parser.add_argument("--output-dir", type=Path, default=Path("artifacts/coreml/surya-ocr-2-coreml-8bit/native_assets"))
+    parser.add_argument("--max-cache-length", type=int, default=512)
+    args = parser.parse_args()
+
+    rt = load_runtime_module(args.runtime_script)
+    output_dir = args.output_dir.expanduser().resolve()
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    processor = AutoProcessor.from_pretrained(args.model_id, trust_remote_code=True)
+    model = rt.load_model(args.model_id, torch.float32)
+    sample = rt.build_sample(processor)
+
+    with torch.no_grad():
+        input_ids = sample["input_ids"].to(torch.long)
+        attention_mask = sample["attention_mask"].to(torch.long)
+        mm_token_type_ids = sample["mm_token_type_ids"].to(torch.long)
+        pixel_values = sample["pixel_values"].to(next(model.parameters()).dtype)
+        image_grid_thw = sample["image_grid_thw"].to(torch.long)
+
+        text_embeds = model.model.get_input_embeddings()(input_ids)
+        image_outputs = model.model.get_image_features(pixel_values, image_grid_thw, return_dict=True)
+        image_embeds = torch.cat(image_outputs.pooler_output, dim=0).to(text_embeds.dtype)
+        image_mask, _ = model.model.get_placeholder_mask(input_ids, inputs_embeds=text_embeds, image_features=image_embeds)
+        image_token_indices = image_mask.squeeze(0).squeeze(-1).nonzero().flatten().to(torch.long)
+        prefill_base = text_embeds.clone()
+        prefill_base[:, image_token_indices, :] = 0
+
+        merged_embeds = text_embeds.clone()
+        merged_embeds[:, image_token_indices, :] = image_embeds.reshape(1, image_embeds.shape[0], image_embeds.shape[1])
+        position_ids = model.model.compute_3d_position_ids(
+            input_ids=input_ids,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=None,
+            inputs_embeds=merged_embeds,
+            attention_mask=attention_mask,
+            past_key_values=None,
+            mm_token_type_ids=mm_token_type_ids,
+        )
+        prefill_cos, prefill_sin = model.model.language_model.rotary_emb(merged_embeds, position_ids)
+
+        one_embed = torch.zeros((1, 1, merged_embeds.shape[-1]), dtype=merged_embeds.dtype)
+        step_cos = []
+        step_sin = []
+        for pos in range(args.max_cache_length):
+            step_position_ids = torch.full((3, 1, 1), pos, dtype=torch.long)
+            cos, sin = model.model.language_model.rotary_emb(one_embed, step_position_ids)
+            step_cos.append(cos)
+            step_sin.append(sin)
+        step_cos = torch.cat(step_cos, dim=1)
+        step_sin = torch.cat(step_sin, dim=1)
+
+        token_embedding = model.model.get_input_embeddings().weight.detach().clone()
+
+    write_fp16(output_dir / "prefill_text_embeds_base_fp16.bin", prefill_base)
+    write_fp16(output_dir / "prefill_cos_fp16.bin", prefill_cos)
+    write_fp16(output_dir / "prefill_sin_fp16.bin", prefill_sin)
+    write_fp16(output_dir / "decode_cos_fp16.bin", step_cos)
+    write_fp16(output_dir / "decode_sin_fp16.bin", step_sin)
+    write_fp16(output_dir / "token_embedding_fp16.bin", token_embedding)
+    write_fp16(output_dir / "canary_pixel_values_fp16.bin", pixel_values)
+
+    constants = {
+        "model_id": args.model_id,
+        "prompt": rt.PROMPT,
+        "input_ids": input_ids.squeeze(0).tolist(),
+        "attention_mask": attention_mask.squeeze(0).tolist(),
+        "mm_token_type_ids": mm_token_type_ids.squeeze(0).tolist(),
+        "image_grid_thw": image_grid_thw.squeeze(0).tolist(),
+        "image_token_indices": image_token_indices.tolist(),
+        "shapes": {
+            "prefill_text_embeds_base": list(prefill_base.shape),
+            "prefill_cos": list(prefill_cos.shape),
+            "prefill_sin": list(prefill_sin.shape),
+            "decode_cos": list(step_cos.shape),
+            "decode_sin": list(step_sin.shape),
+            "token_embedding": list(token_embedding.shape),
+            "pixel_values": list(pixel_values.shape),
+            "image_embeds": list(image_embeds.shape),
+        },
+        "dtype": "float16 little-endian raw binaries",
+        "files": {
+            "prefill_text_embeds_base": "prefill_text_embeds_base_fp16.bin",
+            "prefill_cos": "prefill_cos_fp16.bin",
+            "prefill_sin": "prefill_sin_fp16.bin",
+            "decode_cos": "decode_cos_fp16.bin",
+            "decode_sin": "decode_sin_fp16.bin",
+            "token_embedding": "token_embedding_fp16.bin",
+            "canary_pixel_values": "canary_pixel_values_fp16.bin",
+        },
+    }
+    (output_dir / "surya_native_constants.json").write_text(json.dumps(constants, indent=2) + "\n", encoding="utf-8")
+    vocab = processor.tokenizer.get_vocab()
+    id_to_token = [""] * (max(vocab.values()) + 1)
+    for token, idx in vocab.items():
+        id_to_token[idx] = token
+    (output_dir / "id_to_token.json").write_text(json.dumps(id_to_token, ensure_ascii=False) + "\n", encoding="utf-8")
+    print(json.dumps({"output_dir": str(output_dir), "files": sorted(p.name for p in output_dir.iterdir())}, indent=2))
+
+
+if __name__ == "__main__":
+    main()

validation/native_swift_int8_smoke.json

@@ -0,0 +1,8 @@
+{
+  "runtime": "SwiftPM SuryaCoreMLRuntime",
+  "command": "surya-coreml-smoke --model-dir artifacts/coreml/surya-ocr-2-coreml-8bit --image canary_invoice.png --max-tokens 8 --vision int8",
+  "tokens": [1039, 2009, 2046, 2054, 2047, 2041, 2035, 2037, 1979],
+  "text": "<p>Invoice ",
+  "matches_python_coreml_canary": true,
+  "notes": "Native Swift image preprocessing, INT8 vision, FP16 prefill, FP16 decode step, native token embedding lookup, native KV cache update, and native token decoding."
+}