MLX talker backend + SpeechDecoder overflow chunking

Makefile

@@ -1,4 +1,4 @@
-.PHONY: setup setup-huggingface-cli setup-model-repo download-models download-model download-speakerkit-models build build-cli test \
+.PHONY: setup setup-huggingface-cli setup-model-repo download-models download-model download-speakerkit-models build build-cli graft-mlx-metallib test \
  		clean-package-caches list-devices benchmark-connected-devices benchmark-device benchmark-devices \
 		extract-xcresult build-local-server generate-server generate-server-spec generate-server-code
 
@@ -172,6 +172,20 @@ build-cli:
 	@swift build -c release --product argmax-cli
 
 
+# Only needed to RUN `argmax-cli tts --code-decoder-backend mlx`: command-line
+# SwiftPM cannot compile mlx-swift's Metal shaders (the binary fails at runtime
+# with "Failed to load the default metallib"), while xcodebuild compiles them
+# into the mlx-swift_Cmlx.bundle. Build the bundle once via xcodebuild and
+# graft it next to the SwiftPM release binary. The default coreml backend does
+# not need this.
+graft-mlx-metallib:
+	@echo "Building mlx-swift Metal shader bundle via xcodebuild..."
+	@xcodebuild build -scheme argmax-cli -destination platform=macOS -derivedDataPath .build/xcode -quiet
+	@mkdir -p .build/arm64-apple-macosx/release
+	@cp -R .build/xcode/Build/Products/Debug/mlx-swift_Cmlx.bundle .build/arm64-apple-macosx/release/
+	@echo "Grafted mlx-swift_Cmlx.bundle into .build/arm64-apple-macosx/release/"
+
+
 test:
 	@echo "Running tests..."
 	@swift test -v

Package.swift

@@ -8,7 +8,7 @@ let package = Package(
     name: "argmax-oss-swift",
     platforms: [
         .iOS(.v16),
-        .macOS(.v13),
+        .macOS(.v14),
         .watchOS(.v10),
         .visionOS(.v1)
     ],
@@ -29,6 +29,10 @@ let package = Package(
             name: "SpeakerKit",
             targets: ["SpeakerKit"]
         ),
+        .library(
+            name: "TTSKitMLX",
+            targets: ["TTSKitMLX"]
+        ),
         .executable(
             name: "argmax-cli",
             targets: ["ArgmaxCLI"]
@@ -40,6 +44,9 @@ let package = Package(
     ],
     dependencies: [
         .package(url: "https://github.com/apple/swift-argument-parser.git", from: "1.3.0"),
+        // 0.31.5+ requires swift-tools 6.3; 0.31.4 is the newest release that
+        // builds with the Xcode 26.0 toolchain (Swift 6.2).
+        .package(url: "https://github.com/ml-explore/mlx-swift", exact: "0.31.4"),
     ] + (isServerEnabled() ? [
         .package(url: "https://github.com/vapor/vapor.git", from: "4.115.1"),
         .package(url: "https://github.com/apple/swift-openapi-generator", from: "1.10.2"),
@@ -84,6 +91,19 @@ let package = Package(
             ],
             swiftSettings: swiftSettings()
         ),
+        // MLX talker backend for TTSKit voice cloning. The MLX products are
+        // macOS-only (mlx-swift has no watchOS support); on other platforms
+        // the target compiles empty behind `#if canImport(MLX)`.
+        .target(
+            name: "TTSKitMLX",
+            dependencies: [
+                "TTSKit",
+                .product(name: "MLX", package: "mlx-swift", condition: .when(platforms: [.macOS])),
+                .product(name: "MLXNN", package: "mlx-swift", condition: .when(platforms: [.macOS])),
+                .product(name: "MLXFast", package: "mlx-swift", condition: .when(platforms: [.macOS])),
+            ],
+            swiftSettings: swiftSettings()
+        ),
         .testTarget(
             name: "ArgmaxCoreTests",
             dependencies: [
@@ -124,12 +144,20 @@ let package = Package(
             ],
             swiftSettings: swiftSettings()
         ),
+        .testTarget(
+            name: "TTSKitMLXTests",
+            dependencies: [
+                .target(name: "TTSKitMLX", condition: .when(platforms: [.macOS])),
+            ],
+            swiftSettings: swiftSettings()
+        ),
         .executableTarget(
             name: "ArgmaxCLI",
             dependencies: [
                 "WhisperKit",
                 "TTSKit",
                 "SpeakerKit",
+                .target(name: "TTSKitMLX", condition: .when(platforms: [.macOS])),
                 .product(name: "ArgumentParser", package: "swift-argument-parser"),
             ] + (isServerEnabled() ? [
                 .product(name: "Vapor", package: "vapor"),

Package@swift-6.2.swift

@@ -8,7 +8,7 @@ let package = Package(
     name: "argmax-oss-swift",
     platforms: [
         .iOS(.v16),
-        .macOS(.v13),
+        .macOS(.v14),
         .watchOS(.v10),
         .visionOS(.v1)
     ],
@@ -29,6 +29,10 @@ let package = Package(
             name: "SpeakerKit",
             targets: ["SpeakerKit"]
         ),
+        .library(
+            name: "TTSKitMLX",
+            targets: ["TTSKitMLX"]
+        ),
         .library(
             name: "ArgmaxOSSDynamic",
             type: .dynamic,
@@ -45,6 +49,9 @@ let package = Package(
     ],
     dependencies: [
         .package(url: "https://github.com/apple/swift-argument-parser.git", from: "1.7.0"),
+        // 0.31.5+ requires swift-tools 6.3; 0.31.4 is the newest release that
+        // builds with the Xcode 26.0 toolchain (Swift 6.2).
+        .package(url: "https://github.com/ml-explore/mlx-swift", exact: "0.31.4"),
     ] + (isServerEnabled() ? [
         .package(url: "https://github.com/vapor/vapor.git", from: "4.115.1"),
         .package(url: "https://github.com/apple/swift-openapi-generator", from: "1.10.2"),
@@ -88,6 +95,21 @@ let package = Package(
             ],
             swiftSettings: swiftSettings()
         ),
+        // MLX talker backend for TTSKit voice cloning. The MLX products are
+        // macOS-only (mlx-swift has no watchOS support); on other platforms
+        // the target compiles empty behind `#if canImport(MLX)`. No library
+        // evolution: mlx-swift modules are not resilient, so this target is
+        // excluded from the binary-framework (xcframework) surface.
+        .target(
+            name: "TTSKitMLX",
+            dependencies: [
+                "TTSKit",
+                .product(name: "MLX", package: "mlx-swift", condition: .when(platforms: [.macOS])),
+                .product(name: "MLXNN", package: "mlx-swift", condition: .when(platforms: [.macOS])),
+                .product(name: "MLXFast", package: "mlx-swift", condition: .when(platforms: [.macOS])),
+            ],
+            swiftSettings: swiftSettings(libraryEvolution: false)
+        ),
         .testTarget(
             name: "ArgmaxCoreTests",
             dependencies: [
@@ -128,12 +150,20 @@ let package = Package(
             ],
             swiftSettings: swiftSettings(libraryEvolution: false)
         ),
+        .testTarget(
+            name: "TTSKitMLXTests",
+            dependencies: [
+                .target(name: "TTSKitMLX", condition: .when(platforms: [.macOS])),
+            ],
+            swiftSettings: swiftSettings(libraryEvolution: false)
+        ),
         .executableTarget(
             name: "ArgmaxCLI",
             dependencies: [
                 "WhisperKit",
                 "TTSKit",
                 "SpeakerKit",
+                .target(name: "TTSKitMLX", condition: .when(platforms: [.macOS])),
                 .product(name: "ArgumentParser", package: "swift-argument-parser"),
             ] + (isServerEnabled() ? [
                 .product(name: "Vapor", package: "vapor"),

Sources/ArgmaxCLI/TTSCLI.swift

@@ -5,6 +5,9 @@ import ArgumentParser
 import CoreML
 import Foundation
 import TTSKit
+#if canImport(TTSKitMLX)
+import TTSKitMLX
+#endif
 import WhisperKit
 
 // MARK: - CLI-only conformances for ArgumentParser
@@ -142,6 +145,17 @@ struct TTSCLI: AsyncParsableCommand {
     @Option(name: .long, help: "SpeechDecoder mode: latencyOptimized (lowest time-to-first-audio, 1 frame/call), throughputOptimized (higher throughput, ~4x larger pre-buffer, 4 frames/call), or singleFunction (single-function assets, e.g. the base-family speech decoders). Defaults to latencyOptimized, or singleFunction for -base model presets.")
     var speechDecoderMode: Qwen3SpeechDecoderMode?
 
+    // MARK: - CodeDecoder backend
+
+    @Option(name: .long, help: "CodeDecoder (talker) backend: coreml (default) | mlx — MLX talker with batched ICL prefill and no KV cap; macOS 14+, requires the Base-family mlx-community checkpoint in the local HF cache")
+    var codeDecoderBackend: String = "coreml"
+
+    @Option(name: .long, help: "Qwen3-TTS MLX checkpoint snapshot directory for the talker (default: the cached HF snapshot of the Base-family mlx-community repo)")
+    var mlxModelDir: String?
+
+    @Option(name: .long, help: "MLX talker KV budget in positions (prompt + generated frames)")
+    var mlxMaxSequenceLength: Int = 1024
+
     // MARK: - Compute unit options
 
     @Option(name: .long, help: "Compute units for embedders (TextProjector, CodeEmbedder, MultiCodeEmbedder) {all,cpuOnly,cpuAndGPU,cpuAndNeuralEngine}")
@@ -238,8 +252,37 @@ struct TTSCLI: AsyncParsableCommand {
             config.downloadAdditionalPatterns += config.voiceCloneDownloadPatterns
         }
 
+        // Swap the CoreML CodeDecoder (talker) for the MLX backend if requested.
+        switch codeDecoderBackend {
+            case "coreml":
+                break
+            case "mlx":
+                #if canImport(TTSKitMLX)
+                config.codeDecoder = try makeMlxCodeDecoder(
+                    modelDirectory: mlxModelDir.map { URL(fileURLWithPath: FileManager.resolveAbsolutePath($0)) },
+                    maxSequenceLength: mlxMaxSequenceLength
+                )
+                #else
+                throw ValidationError("--code-decoder-backend mlx is not available on this platform (requires macOS 14+ with MLX support)")
+                #endif
+            default:
+                throw ValidationError("Unknown --code-decoder-backend '\(codeDecoderBackend)' (expected coreml or mlx)")
+        }
+
         // Default: --play uses sequential (1), file output uses unlimited (0).
-        let effectiveWorkerCount = concurrentWorkerCount ?? (play ? 1 : 0)
+        var effectiveWorkerCount = concurrentWorkerCount ?? (play ? 1 : 0)
+
+        // The MLX talker keeps one private KV cache per decoder instance, so
+        // concurrent chunk workers interleave prefills and corrupt it
+        // (broadcast_shapes / retain-count crashes). Serialize generation on
+        // the mlx backend — this also matches the Python prototype, which
+        // synthesizes chunks sequentially.
+        if codeDecoderBackend == "mlx", effectiveWorkerCount != 1 {
+            if concurrentWorkerCount != nil {
+                print("Warning: --code-decoder-backend mlx supports sequential generation only; forcing --concurrent-worker-count 1.")
+            }
+            effectiveWorkerCount = 1
+        }
 
         // Always use a seed for reproducibility -- generate one if not provided
         let effectiveSeed = seed ?? UInt64.random(in: 0...UInt64(UInt32.max))

Sources/TTSKit/Protocols.swift

@@ -44,6 +44,25 @@ public protocol CodeDecoding: MLModelLoading {
     func makeState() -> Any?
 }
 
+// MARK: - Batch Prefill Capability
+
+/// Optional `CodeDecoding` capability: consume an entire prompt prefix in one
+/// forward pass instead of the default token-by-token prefill loop.
+///
+/// Backends with genuinely batched attention (e.g. the MLX code decoder in
+/// `Extensions/TTSKitMLX`) conform to this; the fixed-shape CoreML decoders do
+/// not (their `input_embeds` input is compiled for a single position).
+/// `Qwen3GenerateTask.prefillCodeDecoder` checks for the capability and uses
+/// one batched call when available.
+public protocol BatchPrefillCapable {
+    /// Prefill the full prompt prefix (oldest position first) in one call.
+    ///
+    /// Advances `cache` bookkeeping by `embeds.count` positions and returns the
+    /// decoder output for the last position, exactly as if the prefix had been
+    /// fed through `decode` one position at a time.
+    func prefill(embeds: [[FloatType]], cache: KVCache, state: Any?) async throws -> CodeDecoderOutput
+}
+
 // MARK: - Code Decoder Output
 
 public struct CodeDecoderOutput {
@@ -57,6 +76,22 @@ public struct CodeDecoderOutput {
     /// Time spent on KV cache update inside the decoder (async path only). Lets callers
     /// subtract this from total decode time to isolate pure prediction cost.
     public var internalCacheUpdateTime: TimeInterval = 0
+
+    /// Public memberwise initializer so external `CodeDecoding` implementations
+    /// (e.g. the MLX code decoder in `Extensions/TTSKitMLX`) can construct outputs.
+    public init(
+        logits: any EmbedTensorType,
+        hiddenStates: any EmbedTensorType,
+        keyCacheUpdates: MLMultiArray?,
+        valueCacheUpdates: MLMultiArray?,
+        internalCacheUpdateTime: TimeInterval = 0
+    ) {
+        self.logits = logits
+        self.hiddenStates = hiddenStates
+        self.keyCacheUpdates = keyCacheUpdates
+        self.valueCacheUpdates = valueCacheUpdates
+        self.internalCacheUpdateTime = internalCacheUpdateTime
+    }
 }
 
 // MARK: - Multi-Code Decoder Output

Sources/TTSKit/Qwen3TTS/Qwen3GenerateTask.swift

@@ -51,10 +51,15 @@ open class Qwen3GenerateTask: @unchecked Sendable, SpeechGenerating {
     /// Model components - concrete Qwen3 types for correct async method dispatch.
     /// Using `any Protocol` existentials would cause async extension methods to dispatch
     /// to the protocol default (sync path) instead of the Qwen3-specific MLTensor path.
+    ///
+    /// `codeDecoder` is the exception: everything the task calls on it is a
+    /// `CodeDecoding` protocol requirement (no extension methods), so it stays
+    /// protocol-typed — this is the seam that lets `Extensions/TTSKitMLX`
+    /// substitute an MLX-backed talker via `TTSKitConfig.codeDecoder`.
     public let textProjector: Qwen3TextProjector
     public let codeEmbedder: Qwen3CodeEmbedder
     public let multiCodeEmbedder: Qwen3MultiCodeEmbedder
-    public let codeDecoder: Qwen3CodeDecoder
+    public let codeDecoder: any CodeDecoding
     public let multiCodeDecoder: Qwen3MultiCodeDecoder
     public let speechDecoder: Qwen3SpeechDecoder
     public let sampler: any TokenSampling
@@ -75,7 +80,7 @@ open class Qwen3GenerateTask: @unchecked Sendable, SpeechGenerating {
         textProjector: Qwen3TextProjector,
         codeEmbedder: Qwen3CodeEmbedder,
         multiCodeEmbedder: Qwen3MultiCodeEmbedder,
-        codeDecoder: Qwen3CodeDecoder,
+        codeDecoder: any CodeDecoding,
         multiCodeDecoder: Qwen3MultiCodeDecoder,
         speechDecoder: Qwen3SpeechDecoder,
         sampler: any TokenSampling,
@@ -313,8 +318,15 @@ open class Qwen3GenerateTask: @unchecked Sendable, SpeechGenerating {
 
             // TODO: Remove forking logic with package with min os version upgrade
             if #available(macOS 15.0, iOS 18.0, watchOS 11.0, visionOS 2.0, *), !options.forceLegacyEmbedPath {
-                for embed in combinedEmbeds {
-                    lastCdOutput = try await codeDecoder.decode(inputEmbeds: embed.asMLTensor(), cache: cdCache, state: cdState)
+                if let batchDecoder = codeDecoder as? BatchPrefillCapable {
+                    // Batched prefill: one forward pass over the whole prefix.
+                    // The sequential loop below is dominated by per-call
+                    // dispatch overhead on decoders that support batching.
+                    lastCdOutput = try await batchDecoder.prefill(embeds: combinedEmbeds, cache: cdCache, state: cdState)
+                } else {
+                    for embed in combinedEmbeds {
+                        lastCdOutput = try await codeDecoder.decode(inputEmbeds: embed.asMLTensor(), cache: cdCache, state: cdState)
+                    }
                 }
             } else {
                 for (embedIndex, embed) in combinedEmbeds.enumerated() {