Expose cooperative generation and Core ML placement

Sources/NemotronStreamingASR/NemotronStreamingASR.swift

@@ -191,6 +191,7 @@ public class NemotronStreamingASRModel {
 
     public static func fromPretrained(
         modelId: String? = nil,
+        computeUnits: MLComputeUnits = .all,
         progressHandler: ((Double, String) -> Void)? = nil
     ) async throws -> NemotronStreamingASRModel {
         let effectiveModelId = modelId ?? defaultModelId
@@ -229,23 +230,35 @@ public class NemotronStreamingASRModel {
                 modelId: effectiveModelId, reason: "Download failed", underlying: error)
         }
 
-        return try await load(from: cacheDir, source: effectiveModelId, progressHandler: progressHandler)
+        return try await load(
+            from: cacheDir,
+            source: effectiveModelId,
+            computeUnits: computeUnits,
+            progressHandler: progressHandler
+        )
     }
 
     /// Load a model from a local directory (no download). The directory must
     /// contain `encoder.mlmodelc/`, `decoder.mlmodelc/`, `joint.mlmodelc/`,
     /// `vocab.json`, `languages.json`, and optionally `config.json`.
     public static func fromLocal(
         bundleDir: URL,
+        computeUnits: MLComputeUnits = .all,
         progressHandler: ((Double, String) -> Void)? = nil
     ) async throws -> NemotronStreamingASRModel {
         AudioLog.modelLoading.info("Loading Nemotron Streaming from local: \(bundleDir.path)")
-        return try await load(from: bundleDir, source: bundleDir.path, progressHandler: progressHandler)
+        return try await load(
+            from: bundleDir,
+            source: bundleDir.path,
+            computeUnits: computeUnits,
+            progressHandler: progressHandler
+        )
     }
 
     private static func load(
         from cacheDir: URL,
         source: String,
+        computeUnits: MLComputeUnits,
         progressHandler: ((Double, String) -> Void)?
     ) async throws -> NemotronStreamingASRModel {
         progressHandler?(0.70, "Loading configuration...")
@@ -290,16 +303,19 @@ public class NemotronStreamingASRModel {
             )
         }
 
-        // `.all` lets CoreML schedule the encoder onto the ANE (which is what
-        // Python coremltools' `ComputeUnit.ALL` does). Encoder gains ~40% RTF
-        // over `.cpuAndGPU`. Decoder + joint are tiny enough that ANE vs CPU
-        // is a wash, but using `.all` keeps the unit selection consistent.
+        // Callers that share the GPU with another resident model can exclude
+        // it explicitly with `.cpuAndNeuralEngine`. `.all` remains the generic
+        // default to preserve existing clients and because a narrower placement
+        // still needs a model-, device-, and language-specific parity gate.
         progressHandler?(0.80, "Loading CoreML models...")
-        let encoder = try loadCoreMLModel(name: "encoder", from: cacheDir, computeUnits: .all)
+        let encoder = try loadCoreMLModel(
+            name: "encoder", from: cacheDir, computeUnits: computeUnits)
         progressHandler?(0.90, "Loading decoder...")
-        let decoder = try loadCoreMLModel(name: "decoder", from: cacheDir, computeUnits: .all)
+        let decoder = try loadCoreMLModel(
+            name: "decoder", from: cacheDir, computeUnits: computeUnits)
         progressHandler?(0.95, "Loading joint network...")
-        let joint = try loadCoreMLModel(name: "joint", from: cacheDir, computeUnits: .all)
+        let joint = try loadCoreMLModel(
+            name: "joint", from: cacheDir, computeUnits: computeUnits)
 
         progressHandler?(1.0, "Model loaded")
         AudioLog.modelLoading.info(

Sources/Qwen3Chat/Gemma4Chat.swift

@@ -92,14 +92,33 @@ public final class Gemma4Chat: @unchecked Sendable {
     /// Streaming generation. Suppresses the reasoning channel and only yields answer text.
     public func generateStream(
         messages: [ChatMessage], sampling: ChatSamplingConfig = .default
+    ) -> AsyncThrowingStream<String, Error> {
+        generateStream(
+            messages: messages,
+            sampling: sampling,
+            shouldContinue: { true })
+    }
+
+    /// Streaming generation with cooperative token-boundary cancellation.
+    ///
+    /// MLX evaluation of one token and the initial prompt prefill are atomic,
+    /// but the caller can stop before the next token is scheduled. Returning
+    /// from `decode` also guarantees the producer is finished before a shared
+    /// model is used by the next request.
+    public func generateStream(
+        messages: [ChatMessage],
+        sampling: ChatSamplingConfig = .default,
+        shouldContinue: @escaping @Sendable () -> Bool
     ) -> AsyncThrowingStream<String, Error> {
         AsyncThrowingStream { continuation in
             Task {
                 let promptTokens = Gemma4ChatTemplate.encode(
                     messages: messages, tokenizer: self.gemmaTokenizer)
-                self.decode(promptTokens: promptTokens, sampling: sampling) { text in
-                    continuation.yield(text)
-                }
+                self.decode(
+                    promptTokens: promptTokens,
+                    sampling: sampling,
+                    shouldContinue: shouldContinue,
+                    onText: { text in continuation.yield(text) })
                 continuation.finish()
             }
         }
@@ -118,6 +137,7 @@ public final class Gemma4Chat: @unchecked Sendable {
     func decode(
         promptTokens: [Int],
         sampling: ChatSamplingConfig,
+        shouldContinue: () -> Bool = { true },
         onToken: (Int) -> Void = { _ in },
         onText: (String) -> Void
     ) {
@@ -134,7 +154,7 @@ public final class Gemma4Chat: @unchecked Sendable {
         let endTokens = Array(gemmaTokenizer.eosTokenIds)
 
         var remaining = sampling.maxTokens
-        while remaining > 0 {
+        while remaining > 0 && shouldContinue() {
             remaining -= 1
 
             let next = ChatSampler.sampleOnDevice(

Tests/NemotronStreamingASRTests/NemotronStreamingASRTests.swift

@@ -480,3 +480,78 @@ final class E2ENemotronStreamingASRTests: XCTestCase {
             "English-only bundle should recover every content word; got \(matched)/\(expected)")
     }
 }
+
+/// Opt-in placement gate for Stenograf's Core ML preview candidate. Keeping it
+/// separate from the shared E2E model avoids retaining `.all` and CPU+ANE
+/// copies at once and makes the timing comparison meaningful.
+final class E2ENemotronComputePlacementTests: XCTestCase {
+    func testCPUAndNeuralEngineMatchesAllStreamingOutput() async throws {
+        guard ProcessInfo.processInfo.environment[
+            "NEMOTRON_COMPUTE_PLACEMENT_E2E"
+        ] == "1" else {
+            throw XCTSkip("set NEMOTRON_COMPUTE_PLACEMENT_E2E=1")
+        }
+        let audioURL = Bundle.module.url(
+            forResource: "test_audio", withExtension: "wav")!
+        let audio = try AudioFileLoader.load(
+            url: audioURL, targetSampleRate: 16_000)
+
+        let baseline = try await load(computeUnits: .all)
+        try baseline.warmUp()
+        let baselineStarted = Date()
+        let baselineText = try streamingText(model: baseline, audio: audio)
+        let baselineMilliseconds =
+            Date().timeIntervalSince(baselineStarted) * 1_000
+        baseline.unload()
+
+        let candidate = try await load(computeUnits: .cpuAndNeuralEngine)
+        try candidate.warmUp()
+        let candidateStarted = Date()
+        let candidateText = try streamingText(model: candidate, audio: audio)
+        let candidateMilliseconds =
+            Date().timeIntervalSince(candidateStarted) * 1_000
+
+        XCTAssertEqual(candidateText, baselineText)
+        XCTAssertFalse(candidateText.isEmpty)
+        print(String(
+            format:
+                "[NEMOTRON-PLACEMENT] all=%.2fms cpu+ane=%.2fms parity=%@",
+            baselineMilliseconds,
+            candidateMilliseconds,
+            candidateText == baselineText ? "yes" : "no"))
+    }
+
+    private func load(
+        computeUnits: MLComputeUnits
+    ) async throws -> NemotronStreamingASRModel {
+        if let local = localBundlePath() {
+            return try await NemotronStreamingASRModel.fromLocal(
+                bundleDir: local, computeUnits: computeUnits)
+        }
+        return try await NemotronStreamingASRModel.fromPretrained(
+            computeUnits: computeUnits)
+    }
+
+    private func streamingText(
+        model: NemotronStreamingASRModel,
+        audio: [Float]
+    ) throws -> String {
+        let session = try model.createSession(language: "en-US")
+        let chunkSamples =
+            model.config.streaming.chunkMs * model.config.sampleRate / 1_000
+        var last = ""
+        var cursor = 0
+        while cursor < audio.count {
+            let upper = min(audio.count, cursor + chunkSamples)
+            for partial in try session.pushAudio(Array(audio[cursor..<upper]))
+            where !partial.text.isEmpty {
+                last = partial.text
+            }
+            cursor = upper
+        }
+        for partial in try session.finalize() where !partial.text.isEmpty {
+            last = partial.text
+        }
+        return last
+    }
+}

Tests/Qwen3ChatTests/E2EGemma4GenTests.swift

@@ -47,6 +47,50 @@ final class E2EGemma4GenTests: XCTestCase {
         }
     }
 
+    func testCooperativeCancellationStopsBeforeTheNextToken() throws {
+        guard FileManager.default.fileExists(
+            atPath: Self.modelDir.appendingPathComponent("config.json").path) else {
+            throw XCTSkip("Gemma 4 model dir unavailable: \(Self.modelDir.path)")
+        }
+
+        let chat: Gemma4Chat
+        do {
+            chat = try Gemma4Chat.fromDirectory(Self.modelDir)
+        } catch {
+            throw XCTSkip("model load failed (weights/metallib): \(error)")
+        }
+
+        let prompt = Gemma4ChatTemplate.encode(
+            messages: [
+                ChatMessage(role: .system, content: "Answer directly."),
+                ChatMessage(role: .user, content: "Count upward forever."),
+            ],
+            tokenizer: chat.gemmaTokenizer
+        )
+        let sampling = ChatSamplingConfig(
+            temperature: 0,
+            topK: 0,
+            topP: 1.0,
+            maxTokens: 64,
+            repetitionPenalty: 1.0
+        )
+        var admissionChecks = 0
+        var emittedTokens = 0
+        chat.decode(
+            promptTokens: prompt,
+            sampling: sampling,
+            shouldContinue: {
+                admissionChecks += 1
+                return admissionChecks <= 2
+            },
+            onToken: { _ in emittedTokens += 1 },
+            onText: { _ in }
+        )
+
+        XCTAssertEqual(admissionChecks, 3)
+        XCTAssertEqual(emittedTokens, 2)
+    }
+
     /// Deterministic (no model): the reasoning-channel filter drops a `<|channel>thought … <channel|>`
     /// block and emits only the answer text after it, decoding the SentencePiece byte-fallback tokens.
     func testAnswerFilterSuppressesThoughtChannel() throws {