From 247f98eba4834d2228d0c0beeb448a29ff6be362 Mon Sep 17 00:00:00 2001 From: ivan-digital <42473865+ivan-digital@users.noreply.github.com> Date: Fri, 21 Aug 2026 18:02:18 +0200 Subject: [PATCH] Expose cooperative generation and Core ML placement --- .../NemotronStreamingASR.swift | 34 ++++++--- Sources/Qwen3Chat/Gemma4Chat.swift | 28 ++++++- .../NemotronStreamingASRTests.swift | 75 +++++++++++++++++++ Tests/Qwen3ChatTests/E2EGemma4GenTests.swift | 44 +++++++++++ 4 files changed, 168 insertions(+), 13 deletions(-) diff --git a/Sources/NemotronStreamingASR/NemotronStreamingASR.swift b/Sources/NemotronStreamingASR/NemotronStreamingASR.swift index a6c6063b..15fa2bef 100644 --- a/Sources/NemotronStreamingASR/NemotronStreamingASR.swift +++ b/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,7 +230,12 @@ 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 @@ -237,15 +243,22 @@ public class NemotronStreamingASRModel { /// `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( diff --git a/Sources/Qwen3Chat/Gemma4Chat.swift b/Sources/Qwen3Chat/Gemma4Chat.swift index 437ef0e4..3dc41c5a 100644 --- a/Sources/Qwen3Chat/Gemma4Chat.swift +++ b/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 { + 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 { 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( diff --git a/Tests/NemotronStreamingASRTests/NemotronStreamingASRTests.swift b/Tests/NemotronStreamingASRTests/NemotronStreamingASRTests.swift index ee759212..c0f4bb47 100644 --- a/Tests/NemotronStreamingASRTests/NemotronStreamingASRTests.swift +++ b/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..thought … ` /// block and emits only the answer text after it, decoding the SentencePiece byte-fallback tokens. func testAnswerFilterSuppressesThoughtChannel() throws {