Workaround for bad ML state on MacOS

PlaprePico.mlpackage/Data/com.apple.CoreML/model.mlmodel

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a70122791826c020dc3a1ee6bfadef2a5ac74d14e5e060e9dfab76c57284130b
+oid sha256:b0ea4fbe5939f8db381da0ccadf9e90b61c82f5f0eca58b46e89b3a5541a49f0
 size 957824

PlaprePico.mlpackage/Manifest.json

@@ -1,18 +1,18 @@
 {
     "fileFormatVersion": "1.0.0",
     "itemInfoEntries": {
-        "9B90755A-51A8-4710-8C16-AB5B86538A1A": {
-            "author": "com.apple.CoreML",
-            "description": "CoreML Model Specification",
-            "name": "model.mlmodel",
-            "path": "com.apple.CoreML/model.mlmodel"
-        },
-        "D85A4A62-BBF3-4C94-848E-7E37C3571EA6": {
+        "1F911078-42FE-4F91-A2D0-E5B86F87F7AD": {
             "author": "com.apple.CoreML",
             "description": "CoreML Model Weights",
             "name": "weights",
             "path": "com.apple.CoreML/weights"
+        },
+        "3E69D1BF-E09D-43D9-A7FE-E3B15CDDF0BD": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
         }
     },
-    "rootModelIdentifier": "9B90755A-51A8-4710-8C16-AB5B86538A1A"
+    "rootModelIdentifier": "3E69D1BF-E09D-43D9-A7FE-E3B15CDDF0BD"
 }

scripts/convert.py

@@ -207,7 +207,7 @@ def convert_decode(model: PlaprePico, output_dir: Path):
         ],
         outputs=[ct.TensorType(name="logits", dtype=np.float16)],
         states=build_kv_cache_states(),
-        compute_precision=ct.precision.FLOAT32,
+        compute_precision=ct.precision.FLOAT16,
         minimum_deployment_target=ct.target.iOS18,
     )
 

swift-cli/Sources/main.swift

@@ -378,56 +378,9 @@ func sampleFromLogitsFp16(_ ptr: UnsafeBufferPointer<Float16>, temperature: Floa
     return Int32(topIndices[topK - 1])
 }
 
-func sampleFromLogitsFp32(_ ptr: UnsafeBufferPointer<Float>, temperature: Float, topK: Int) -> Int32 {
-    var topIndices = [Int](repeating: 0, count: topK)
-    var topValues = [Float](repeating: -.greatestFiniteMagnitude, count: topK)
-    var minIdx = 0
-    for i in 0..<vocabSize {
-        if ptr[i] > topValues[minIdx] {
-            topValues[minIdx] = ptr[i]
-            topIndices[minIdx] = i
-            minIdx = 0
-            for j in 1..<topK { if topValues[j] < topValues[minIdx] { minIdx = j } }
-        }
-    }
-    if temperature <= 0 {
-        var bestIdx = 0
-        for j in 1..<topK { if topValues[j] > topValues[bestIdx] { bestIdx = j } }
-        return Int32(topIndices[bestIdx])
-    }
-    var logits32 = [Float](repeating: 0, count: topK)
-    for j in 0..<topK { logits32[j] = topValues[j] / temperature }
-    let maxVal = logits32.max()!
-    var exps = logits32.map { exp($0 - maxVal) }
-    let sum = exps.reduce(0, +)
-    for j in 0..<topK { exps[j] /= sum }
-    let r = Float.random(in: 0..<1)
-    var cumsum: Float = 0
-    for j in 0..<topK {
-        cumsum += exps[j]
-        if cumsum >= r { return Int32(topIndices[j]) }
-    }
-    return Int32(topIndices[topK - 1])
-}
-
-func sampleFromLogits(_ logitsArr: MLMultiArray, temperature: Float = 0.8, topK: Int = 50) -> Int32 {
-    // CoreML may report .float16 dataType but use float32 backing with FLOAT32 compute precision.
-    // Try fp16 first; if values are NaN (fp32 data read as fp16), fall back to fp32.
-    var isFp16 = true
-    logitsArr.withUnsafeBufferPointer(ofType: Float16.self) { ptr in
-        if ptr[0].isNaN && ptr[1].isNaN { isFp16 = false }
-    }
-    if isFp16 {
-        return logitsArr.withUnsafeBufferPointer(ofType: Float16.self) { ptr -> Int32 in
-            return sampleFromLogitsFp16(ptr, temperature: temperature, topK: topK)
-        }
-    } else {
-        // fp32 backing behind fp16-declared output — use dataPointer directly
-        let rawPtr = UnsafeBufferPointer(
-            start: logitsArr.dataPointer.assumingMemoryBound(to: Float.self),
-            count: vocabSize
-        )
-        return sampleFromLogitsFp32(rawPtr, temperature: temperature, topK: topK)
+func sampleFromLogits(_ logits: [Float16], temperature: Float = 0.8, topK: Int = 50) -> Int32 {
+    return logits.withUnsafeBufferPointer { ptr -> Int32 in
+        return sampleFromLogitsFp16(ptr, temperature: temperature, topK: topK)
     }
 }
 
@@ -510,8 +463,11 @@ let decodeModel = try measure("Compile PlaprePico") { try compileModel(at: model
 
 // === Step 1: Prefill via decode model (one token at a time) ===
 print("\n--- Prefill (token-by-token through decode model) ---")
-let state = decodeModel.makeState()
-var lastLogitsArr: MLMultiArray!
+
+// MLState may contain uninitialized memory (NaN in KV cache).
+// Retry makeState + first prefill step until logits are valid.
+var state = decodeModel.makeState()
+var lastLogits = [Float16](repeating: 0, count: vocabSize)
 
 // Pre-allocate all input arrays ONCE
 let pInputIds = try! MLMultiArray(shape: [1, 1], dataType: .int32)
@@ -575,7 +531,11 @@ func runDecodeStep(token: Int32, pos: Int, isSpeaker: Bool = false) throws {
     }
 
     let output = try decodeModel.prediction(from: inputProvider, using: state)
-    lastLogitsArr = output.featureValue(for: "logits")!.multiArrayValue!
+    let arr = output.featureValue(for: "logits")!.multiArrayValue!
+    // Log shape/stride info on first call to diagnose backing type
+    arr.withUnsafeBufferPointer(ofType: Float16.self) { ptr in
+        for i in 0..<vocabSize { lastLogits[i] = ptr[i] }
+    }
 }
 
 // The input sequence is: [placeholder(speaker), <text>, tokens..., <audio>]
@@ -589,15 +549,33 @@ func runDecodeStep(token: Int32, pos: Int, isSpeaker: Bool = false) throws {
 let inputTokens: [Int32] = Array(inputSeq.prefix(inputLen))
 print("Processing \(inputTokens.count) input tokens...")
 
+// Retry prefill if state has uninitialized NaN memory.
+// MLState buffers are not guaranteed zero-initialized; NaN in KV cache
+// propagates through Q@K^T and poisons softmax. Typically 1-3 attempts.
+var prefillAttempt = 0
 let prefillStart = CFAbsoluteTimeGetCurrent()
-for (i, token) in inputTokens.enumerated() {
-    try runDecodeStep(token: token, pos: i, isSpeaker: i == 0)
+while prefillAttempt < 20 {
+    prefillAttempt += 1
+    state = decodeModel.makeState()
+    // Reset causal mask and update mask for fresh prefill
+    pCausalMask.withUnsafeMutableBufferPointer(ofType: Float16.self) { ptr, _ in
+        for i in 0..<maxContext { ptr[i] = Float16(-65504.0) }
+    }
+    pUpdateMask.withUnsafeMutableBufferPointer(ofType: Float16.self) { ptr, _ in
+        for i in 0..<maxContext { ptr[i] = Float16(0.0) }
+    }
+    try runDecodeStep(token: inputTokens[0], pos: 0, isSpeaker: true)
+    if !lastLogits[0].isNaN { break }
+}
+print("  Clean state after \(prefillAttempt) attempt(s)")
+for i in 1..<inputTokens.count {
+    try runDecodeStep(token: inputTokens[i], pos: i, isSpeaker: false)
 }
 let prefillElapsed = CFAbsoluteTimeGetCurrent() - prefillStart
 let prefillTokPerSec = Double(inputTokens.count) / prefillElapsed
 print("  ⏱ Prefill: \(formatTime(prefillElapsed)) (\(inputTokens.count) tokens, \(String(format: "%.1f", prefillTokPerSec)) tok/s)")
 
-let firstToken = sampleFromLogits(lastLogitsArr, temperature: 0.8, topK: 50)
+let firstToken = sampleFromLogits(lastLogits, temperature: 0.8, topK: 50)
 print("First generated token: \(firstToken)")
 
 // === Step 2: Autoregressive decode ===
@@ -614,7 +592,7 @@ for step in 1..<maxTokens {
     let pos = inputLen + step - 1
 
     try runDecodeStep(token: nextToken, pos: pos)
-    nextToken = sampleFromLogits(lastLogitsArr, temperature: 0.8, topK: 50)
+    nextToken = sampleFromLogits(lastLogits, temperature: 0.8, topK: 50)
     generatedTokens.append(nextToken)
 
     if nextToken == eosToken {