From 2305550c92a2d000b0e2d4517841aa7e986ab173 Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Sun, 24 May 2026 18:05:37 -0400 Subject: [PATCH] perf(asr/nemotron-multilingual): smart-spec default-on with K auto-detect + outputBackings hot path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Smart-spec V2 default-on: opt-out via FLUIDAUDIO_ENABLE_SMART_SPECULATIVE=0 (also accepts "false"/"no"). Missing assets → transparent fallback to legacy inner loop. - K (speculative window width) read from jointNoEncProjBatched's encoder_proj input shape at load time, stored as jointNoEncProjBatchedK. Eliminates the hardcoded K=8; same Swift hot loop ships per-bundle with matched K. - MLPredictionOptions.outputBackings on the jointBatched hot path. Pre- allocated [1, K, 1, V] logits backing reused per call instead of fresh MLMultiArray alloc. - Load-time state log reports resolved smart-spec state (enabled / disabled / fallback). Warning when env-var explicitly set but assets missing. S6 verification A/B/A/B at 1120ms LS test-clean (n=2620, same session): force-off mean 99.36 vs default-on mean 101.44 → +2.10% RTFx, non-overlapping (101.14-101.74 vs 98.86-99.86). WER delta +0.0005pp. Independently reproduces T3's same-session K=4 +2.0% A/B/A/B. --- ...otronMultilingualAsrManager+Pipeline.swift | 23 ++++++- ...eamingNemotronMultilingualAsrManager.swift | 62 +++++++++++++++++++ 2 files changed, 82 insertions(+), 3 deletions(-) diff --git a/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager+Pipeline.swift b/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager+Pipeline.swift index 90c78fde..699ed4c7 100644 --- a/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager+Pipeline.swift +++ b/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager+Pipeline.swift @@ -297,7 +297,19 @@ extension StreamingNemotronMultilingualAsrManager { // Activates when joint_noencproj_batched.mlpackage is loaded // AND either: (A) encoder emits encoder_proj OR (B) native // weights are available for the Swift-side projection. - let smartSpecEnabled = ProcessInfo.processInfo.environment["FLUIDAUDIO_ENABLE_SMART_SPECULATIVE"] != nil + // + // Default-on as of May 2026 (T3 confirmed K=4 at 1120ms is + // +2.0% non-overlapping, K=8 at 4480ms is +1.7% non-overlapping, + // both WER-neutral). Opt-out via FLUIDAUDIO_ENABLE_SMART_SPECULATIVE=0 + // (or "false"). When the required assets aren't shipped, the path + // falls back transparently to the legacy inner loop regardless. + let smartSpecEnabled: Bool + if let v = ProcessInfo.processInfo.environment["FLUIDAUDIO_ENABLE_SMART_SPECULATIVE"] { + let lowered = v.lowercased() + smartSpecEnabled = !(lowered == "0" || lowered == "false" || lowered == "no") + } else { + smartSpecEnabled = true + } let useSwiftEncProj = (encoderProj == nil) && (self.nativeRnnt != nil) if smartSpecEnabled, let jointBatched = self.jointNoEncProjBatched, @@ -951,7 +963,7 @@ extension StreamingNemotronMultilingualAsrManager { newTokens: inout [Int], tokenizer: NemotronMultilingualTokenizer ) async throws { - let K = 8 + let K = self.jointNoEncProjBatchedK let blankIdx = config.blankIdx // Pre-allocate the batched encoder_proj slice buffer [1, K, 640]. @@ -1048,7 +1060,12 @@ extension StreamingNemotronMultilingualAsrManager { "encoder_proj": MLFeatureValue(multiArray: encProjBatchBuf), "decoder": MLFeatureValue(multiArray: decOut), ]) - let jointOutput = try await jointBatched.prediction(from: jointInput) + let jointOutput: MLFeatureProvider + if let opts = jointNoEncProjBatchedPredictionOptions { + jointOutput = try await jointBatched.prediction(from: jointInput, options: opts) + } else { + jointOutput = try await jointBatched.prediction(from: jointInput) + } guard let logits = jointOutput.featureValue(for: "logits")?.multiArrayValue else { throw ASRError.processingFailed("Speculative joint_noencproj_batched failed") } diff --git a/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager.swift b/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager.swift index 32bab08d..80502c34 100644 --- a/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager.swift +++ b/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager.swift @@ -61,6 +61,12 @@ public actor StreamingNemotronMultilingualAsrManager { /// Used by the smarter speculative-blank decode path that fast-skips /// blank-streaks K-at-a-time. internal var jointNoEncProjBatched: MLModel? + /// K (speculative window width) read from jointNoEncProjBatched's + /// encoder_proj input shape at load time. The Swift inner loop must + /// always match the loaded mlpackage's K — hardcoding here would + /// shape-mismatch any non-K=8 build. Defaults to 8 (no-op fallback) + /// when the model isn't present. + internal var jointNoEncProjBatchedK: Int = 8 /// Optional native-Accelerate RNN-T inner-loop replacement. When present, /// the per-token decoder + joint CoreML calls are replaced by pure-Swift /// vDSP/cblas forward (~5-10x fewer dispatch overhead per token). Loaded @@ -121,6 +127,12 @@ public actor StreamingNemotronMultilingualAsrManager { internal var decoderJointPredictionOptions: MLPredictionOptions? internal var decoderJointArgmaxPredictionOptions: MLPredictionOptions? internal var decoderJointNoEncProjPredictionOptions: MLPredictionOptions? + /// Output backings for the smart-spec batched joint hot path. + /// runSpeculativeBlankDecodeV2 calls jointBatched.prediction() once per + /// K-frame window; pre-allocating the [1, K, 1, V] logits backing once + /// avoids per-call MLMultiArray allocation. Only populated when + /// jointNoEncProjBatched is loaded. + internal var jointNoEncProjBatchedPredictionOptions: MLPredictionOptions? /// Reusable per-frame encoder step buffer. Refilled in-place inside the /// inner RNN-T greedy loop instead of allocating a fresh [1, 1024, 1] /// every emitted token. @@ -399,6 +411,18 @@ public actor StreamingNemotronMultilingualAsrManager { self.jointNoEncProjBatched = try await MLModel.load(contentsOf: tempCompiledURL, configuration: mlConfiguration) logger.info("Compiled + loaded joint_noencproj_batched.mlpackage — smart speculative-blank path available") } + // Read K from the loaded model's encoder_proj input shape so the + // Swift hot loop always matches the asset (K=8 historically; K=4 + // build under evaluation at 1120ms). + if let m = self.jointNoEncProjBatched, + let constraint = m.modelDescription.inputDescriptionsByName["encoder_proj"]?.multiArrayConstraint, + constraint.shape.count >= 2 { + let kFromModel = constraint.shape[1].intValue + if kFromModel > 0 { + self.jointNoEncProjBatchedK = kFromModel + logger.info("Smart-spec K = \(kFromModel) (from joint_noencproj_batched encoder_proj input shape)") + } + } // Optional native-Accelerate RNN-T inner loop. Loaded from // `/native_weights/` (weights.bin + weights_index.json). @@ -413,6 +437,43 @@ public actor StreamingNemotronMultilingualAsrManager { } } + // Smart-speculative-blank load-time state report. Smart-spec is + // default-on as of May 2026 (T3 K=4 at 1120ms = +2.0%, K=8 at + // 4480ms = +1.7%, both A/B/A/B non-overlapping, WER-neutral). + // Honor explicit opt-out: env-var = "0"/"false"/"no" disables. + // Missing assets → transparent fallback to legacy inner loop. + let smartSpecEnvVar = ProcessInfo.processInfo.environment["FLUIDAUDIO_ENABLE_SMART_SPECULATIVE"] + let smartSpecExplicitlyDisabled: Bool = { + guard let v = smartSpecEnvVar?.lowercased() else { return false } + return v == "0" || v == "false" || v == "no" + }() + let smartSpecEnabledForLogging = !smartSpecExplicitlyDisabled + var smartSpecMissing: [String] = [] + if self.jointNoEncProjBatched == nil { + smartSpecMissing.append("joint_noencproj_batched.mlpackage") + } + if self.nativeRnnt == nil { + smartSpecMissing.append("native_weights/") + } + if smartSpecExplicitlyDisabled { + logger.info("Smart-spec: explicitly disabled via FLUIDAUDIO_ENABLE_SMART_SPECULATIVE=\(smartSpecEnvVar ?? "")") + } else if smartSpecMissing.isEmpty { + logger.info("Smart-spec: enabled (default-on; assets present; K=\(self.jointNoEncProjBatchedK))") + } else { + // Default-on intent, but assets missing → legacy fallback. + // Warn only if the user explicitly opted IN with the env-var + // (they probably expected smart-spec to run); otherwise emit + // info because the operator may have intentionally trimmed + // the bundle. + let msg = "Smart-spec: assets missing (\(smartSpecMissing.joined(separator: ", "))); falling back to legacy inner loop" + if smartSpecEnvVar != nil { + logger.warning(msg) + } else { + logger.info(msg) + } + } + _ = smartSpecEnabledForLogging // silence unused warning if introspected later + // Load tokenizer with lang-tag filter set let tokenizerURL = directory.appendingPathComponent(ModelNames.NemotronMultilingualStreaming.tokenizer) self.tokenizer = try NemotronMultilingualTokenizer( @@ -432,6 +493,7 @@ public actor StreamingNemotronMultilingualAsrManager { self.decoderJointPredictionOptions = Self.makePredictionOptions(for: self.decoderJoint) self.decoderJointArgmaxPredictionOptions = Self.makePredictionOptions(for: self.decoderJointArgmax) self.decoderJointNoEncProjPredictionOptions = Self.makePredictionOptions(for: self.decoderJointNoEncProj) + self.jointNoEncProjBatchedPredictionOptions = Self.makePredictionOptions(for: self.jointNoEncProjBatched) // Reusable inner-loop step buffers ([1, encoder_dim, 1] and // [1, 1, joint_dim] for the B3 path). self.encoderStepBuf = try? MLMultiArray(shape: [1, NSNumber(value: config.encoderDim), 1], dataType: .float32)