Add Apple MLX backend for native Metal GPU acceleration on macOS

Introduce a dual-backend architecture (tch-backend vs mlx) with a unified
Tensor abstraction, following the pattern from qwen3_tts_rs. The MLX backend
builds mlx-c from a git submodule and uses Metal GPU natively on Apple Silicon,
eliminating the libtorch dependency on macOS.

- Add src/tensor.rs unified Tensor type wrapping tch::Tensor or MlxArray via #[cfg]
- Add src/backend/mlx/ with FFI bindings, RAII array wrapper, ops, I/O, signal processing
- Add build.rs for CMake-based mlx-c compilation
- Port all model code (layers, encoder, decoder, inference) to use unified Tensor
- Update CI to test both tch and mlx backends on macOS
- Update release workflow to build mlx for macOS, tch for Linux
- Fix release zip structure to expand into a named directory
- Update skills/ for macOS MLX (no libtorch needed) and 0.6B model only

Signed-off-by: Michael Yuan <michael@secondstate.io>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: juntao

.github/workflows/ci.yml

@@ -12,22 +12,34 @@ jobs:
         include:
           - os: ubuntu-latest
             name: Linux x86_64
+            backend: tch
             libtorch-url: https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.7.1%2Bcpu.zip
             libtorch-archive: libtorch.zip
           - os: ubuntu-24.04-arm
             name: Linux ARM64
+            backend: tch
             libtorch-url: https://github.com/second-state/libtorch-releases/releases/download/v2.7.1/libtorch-cxx11-abi-aarch64-2.7.1.tar.gz
             libtorch-archive: libtorch.tar.gz
           - os: macos-latest
-            name: macOS ARM64
+            name: macOS ARM64 (tch)
+            backend: tch
             libtorch-url: https://download.pytorch.org/libtorch/cpu/libtorch-macos-arm64-2.7.1.zip
             libtorch-archive: libtorch.zip
+          - os: macos-latest
+            name: macOS ARM64 (mlx)
+            backend: mlx
 
     runs-on: ${{ matrix.os }}
     name: Build (${{ matrix.name }})
 
     steps:
       - uses: actions/checkout@v4
+        if: matrix.backend == 'tch'
+
+      - uses: actions/checkout@v4
+        if: matrix.backend == 'mlx'
+        with:
+          submodules: recursive
 
       - name: Install build tools (Linux)
         if: runner.os == 'Linux'
@@ -42,14 +54,16 @@ jobs:
           path: |
             ~/.cargo/registry
             ~/.cargo/git
-          key: ${{ runner.os }}-${{ runner.arch }}-cargo-registry-${{ hashFiles('**/Cargo.lock') }}
+          key: ${{ runner.os }}-${{ runner.arch }}-${{ matrix.backend }}-cargo-registry-${{ hashFiles('**/Cargo.lock') }}
           restore-keys: |
-            ${{ runner.os }}-${{ runner.arch }}-cargo-registry-
+            ${{ runner.os }}-${{ runner.arch }}-${{ matrix.backend }}-cargo-registry-
 
       - name: Download libtorch
+        if: matrix.backend == 'tch'
         run: curl -Lo ${{ matrix.libtorch-archive }} "${{ matrix.libtorch-url }}"
 
       - name: Extract libtorch
+        if: matrix.backend == 'tch'
         run: |
           if [[ "${{ matrix.libtorch-archive }}" == *.zip ]]; then
             unzip -q ${{ matrix.libtorch-archive }}
@@ -58,21 +72,26 @@ jobs:
           fi
 
       - name: Set linker rpath-link (Linux only)
-        if: runner.os == 'Linux'
+        if: matrix.backend == 'tch' && runner.os == 'Linux'
         run: echo "RUSTFLAGS=-C link-arg=-Wl,-rpath-link,${{ github.workspace }}/libtorch/lib" >> "$GITHUB_ENV"
 
       - name: Cap libtorch CPU ISA to AVX2 (x86_64 only)
-        if: runner.arch == 'X64'
+        if: matrix.backend == 'tch' && runner.arch == 'X64'
         run: |
           echo "DNNL_MAX_CPU_ISA=AVX2" >> "$GITHUB_ENV"
           echo "ATEN_CPU_CAPABILITY=avx2" >> "$GITHUB_ENV"
 
-      - name: Build
+      - name: Build (tch)
+        if: matrix.backend == 'tch'
         env:
           LIBTORCH: ${{ github.workspace }}/libtorch
           LIBTORCH_BYPASS_VERSION_CHECK: "1"
         run: cargo build --release --features build-ffmpeg
 
+      - name: Build (mlx)
+        if: matrix.backend == 'mlx'
+        run: cargo build --release --no-default-features --features mlx,build-ffmpeg
+
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
@@ -82,11 +101,11 @@ jobs:
         run: pip install huggingface_hub transformers
 
       - name: Set library path (Linux)
-        if: runner.os == 'Linux'
+        if: matrix.backend == 'tch' && runner.os == 'Linux'
         run: echo "LD_LIBRARY_PATH=${{ github.workspace }}/libtorch/lib:$LD_LIBRARY_PATH" >> "$GITHUB_ENV"
 
       - name: Set library path (macOS)
-        if: runner.os == 'macOS'
+        if: matrix.backend == 'tch' && runner.os == 'macOS'
         run: echo "DYLD_LIBRARY_PATH=${{ github.workspace }}/libtorch/lib:$DYLD_LIBRARY_PATH" >> "$GITHUB_ENV"
 
       - name: Cache 0.6B model weights
@@ -146,5 +165,5 @@ jobs:
       - name: Upload binary
         uses: actions/upload-artifact@v4
         with:
-          name: qwen3-asr-${{ matrix.os }}-${{ runner.arch }}.zip
+          name: qwen3-asr-${{ matrix.os }}-${{ runner.arch }}-${{ matrix.backend }}.zip
           path: target/release/asr

.github/workflows/release.yml

@@ -12,19 +12,20 @@ jobs:
         include:
           - os: ubuntu-latest
             name: Linux x86_64
+            backend: tch
             asset-name: asr-linux-x86_64
             libtorch-url: https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.7.1%2Bcpu.zip
             libtorch-archive: libtorch.zip
           - os: ubuntu-24.04-arm
             name: Linux ARM64
+            backend: tch
             asset-name: asr-linux-aarch64
             libtorch-url: https://github.com/second-state/libtorch-releases/releases/download/v2.7.1/libtorch-cxx11-abi-aarch64-2.7.1.tar.gz
             libtorch-archive: libtorch.tar.gz
           - os: macos-latest
             name: macOS ARM64
+            backend: mlx
             asset-name: asr-macos-aarch64
-            libtorch-url: https://download.pytorch.org/libtorch/cpu/libtorch-macos-arm64-2.7.1.zip
-            libtorch-archive: libtorch.zip
 
     runs-on: ${{ matrix.os }}
     name: Build (${{ matrix.name }})
@@ -34,6 +35,12 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
+        if: matrix.backend == 'tch'
+
+      - uses: actions/checkout@v4
+        if: matrix.backend == 'mlx'
+        with:
+          submodules: recursive
 
       - name: Install build tools (Linux)
         if: runner.os == 'Linux'
@@ -48,14 +55,16 @@ jobs:
           path: |
             ~/.cargo/registry
             ~/.cargo/git
-          key: ${{ runner.os }}-${{ runner.arch }}-cargo-registry-${{ hashFiles('**/Cargo.lock') }}
+          key: ${{ runner.os }}-${{ runner.arch }}-${{ matrix.backend }}-cargo-registry-${{ hashFiles('**/Cargo.lock') }}
           restore-keys: |
-            ${{ runner.os }}-${{ runner.arch }}-cargo-registry-
+            ${{ runner.os }}-${{ runner.arch }}-${{ matrix.backend }}-cargo-registry-
 
       - name: Download libtorch
+        if: matrix.backend == 'tch'
         run: curl -Lo ${{ matrix.libtorch-archive }} "${{ matrix.libtorch-url }}"
 
       - name: Extract libtorch
+        if: matrix.backend == 'tch'
         run: |
           if [[ "${{ matrix.libtorch-archive }}" == *.zip ]]; then
             unzip -q ${{ matrix.libtorch-archive }}
@@ -64,19 +73,25 @@ jobs:
           fi
 
       - name: Set linker rpath-link (Linux only)
-        if: runner.os == 'Linux'
+        if: matrix.backend == 'tch' && runner.os == 'Linux'
         run: echo "RUSTFLAGS=-C link-arg=-Wl,-rpath-link,${{ github.workspace }}/libtorch/lib" >> "$GITHUB_ENV"
 
-      - name: Build
+      - name: Build (tch)
+        if: matrix.backend == 'tch'
         env:
           LIBTORCH: ${{ github.workspace }}/libtorch
           LIBTORCH_BYPASS_VERSION_CHECK: "1"
         run: cargo build --release --features build-ffmpeg
 
+      - name: Build (mlx)
+        if: matrix.backend == 'mlx'
+        run: cargo build --release --no-default-features --features mlx,build-ffmpeg
+
       - name: Package binary
         run: |
-          cp target/release/asr ${{ matrix.asset-name }}
-          zip ${{ matrix.asset-name }}.zip ${{ matrix.asset-name }}
+          mkdir ${{ matrix.asset-name }}
+          cp target/release/asr ${{ matrix.asset-name }}/asr
+          zip -r ${{ matrix.asset-name }}.zip ${{ matrix.asset-name }}
 
       - name: Upload release asset
         env:

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "mlx-c"]
+	path = mlx-c
+	url = https://github.com/ml-explore/mlx-c.git

Cargo.lock

@@ -10,7 +10,7 @@ name = "asr"
 path = "src/main.rs"
 
 [dependencies]
-tch = "0.20"
+tch = { version = "0.20", optional = true }
 tokenizers = "0.21"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
@@ -24,7 +24,12 @@ safetensors = "0.5"
 ffmpeg-next = "8"
 ffmpeg-sys-next = "8"
 
+[build-dependencies]
+cmake = "0.1"
+
 [features]
-default = []
+default = ["tch-backend"]
+tch-backend = ["dep:tch"]
+mlx = []
 static-ffmpeg = ["ffmpeg-sys-next/static"]
 build-ffmpeg = ["ffmpeg-sys-next/build", "static-ffmpeg"]

Cargo.toml

@@ -1,6 +1,6 @@
 # Qwen3 ASR Rust
 
-Pure Rust implementation of [Qwen3-ASR](https://github.com/QwenLM/Qwen3-ASR) automatic speech recognition using libtorch. Loads model weights directly from safetensors files and re-implements the complete neural network forward pass in Rust.
+Pure Rust implementation of [Qwen3-ASR](https://github.com/QwenLM/Qwen3-ASR) automatic speech recognition. Supports two backends: **libtorch** (via the `tch` crate, cross-platform with optional CUDA) and **MLX** (Apple Silicon native via Metal GPU). Loads model weights directly from safetensors files and re-implements the complete neural network forward pass in Rust.
 
 ## Architecture
 
@@ -19,7 +19,16 @@ The implementation ports the Qwen3-ASR encoder-decoder architecture from PyTorch
 
 ## Prerequisites
 
-### libtorch
+### Backend
+
+Choose one backend:
+
+| Backend | Feature flag | Platforms | GPU |
+|---------|-------------|-----------|-----|
+| libtorch | `tch-backend` (default) | Linux, macOS, Windows | CUDA |
+| MLX | `mlx` | macOS Apple Silicon | Metal |
+
+### libtorch (for `tch-backend`)
 
 The `tch` crate (v0.20) requires **libtorch 2.7.1**. Download and extract for your platform:
 
@@ -77,6 +86,8 @@ tok.backend_tokenizer.save('Qwen3-ASR-0.6B/tokenizer.json')
 
 ## Build
 
+### libtorch backend (default)
+
 ```bash
 # Set environment
 export LIBTORCH=$(pwd)/libtorch
@@ -94,6 +105,19 @@ cargo build --release --features static-ffmpeg
 cargo build --release --features build-ffmpeg
 ```
 
+### MLX backend (macOS Apple Silicon)
+
+```bash
+# Initialize mlx-c submodule
+git submodule update --init --recursive
+
+# Build with MLX (no libtorch needed)
+cargo build --release --no-default-features --features mlx
+
+# With statically linked FFmpeg
+cargo build --release --no-default-features --features mlx,static-ffmpeg
+```
+
 ## Usage
 
 ```bash
@@ -134,6 +158,7 @@ Qwen3-ASR supports 30 languages: Chinese, English, Cantonese, Arabic, German, Fr
 src/
 ├── main.rs            # CLI binary entry point
 ├── lib.rs             # Library module declarations
+├── tensor.rs          # Unified Tensor abstraction (tch/MLX backend)
 ├── config.rs          # Model configuration (from config.json)
 ├── error.rs           # Error types
 ├── audio.rs           # FFmpeg-based audio loading and format conversion
@@ -144,7 +169,15 @@ src/
 ├── audio_encoder.rs   # Whisper-style audio encoder (Conv2d + Transformer)
 ├── text_decoder.rs    # Qwen3 text decoder with KV cache
 ├── tokenizer.rs       # HuggingFace tokenizer wrapper
-└── inference.rs       # End-to-end ASR inference pipeline
+├── inference.rs       # End-to-end ASR inference pipeline
+└── backend/
+    └── mlx/           # Apple MLX backend (Metal GPU)
+        ├── ffi.rs     # Raw C FFI bindings to mlx-c
+        ├── array.rs   # Safe RAII MlxArray wrapper
+        ├── ops.rs     # Safe operation wrappers
+        ├── io.rs      # Safetensors loading via mlx-c
+        ├── signal.rs  # STFT, mel spectrogram signal processing
+        └── stream.rs  # Device/stream management
 ```
 
 ## License

README.md

@@ -0,0 +1,72 @@
+// Build script for qwen3_asr.
+//
+// When the `mlx` feature is enabled (macOS only), this script:
+// 1. Builds the mlx-c library from the git submodule via CMake
+// 2. Emits linker directives for mlx-c, MLX, Metal, and system frameworks
+
+fn main() {
+    #[cfg(feature = "mlx")]
+    build_mlx();
+}
+
+#[cfg(feature = "mlx")]
+fn build_mlx() {
+    let target_os = std::env::var("CARGO_CFG_TARGET_OS").unwrap_or_default();
+    let target_arch = std::env::var("CARGO_CFG_TARGET_ARCH").unwrap_or_default();
+
+    if target_os != "macos" {
+        panic!("The `mlx` feature is only supported on macOS. Current target OS: {target_os}");
+    }
+    if target_arch != "aarch64" {
+        eprintln!(
+            "Warning: MLX is optimized for Apple Silicon (aarch64). \
+             Current target arch: {target_arch}. Metal GPU acceleration may not be available."
+        );
+    }
+
+    let mlx_c_dir = std::path::PathBuf::from("mlx-c");
+    if !mlx_c_dir.join("CMakeLists.txt").exists() {
+        panic!(
+            "mlx-c submodule not found. Please run:\n\
+             \n\
+             git submodule update --init --recursive\n\
+             \n\
+             to clone the mlx-c dependency."
+        );
+    }
+
+    // Build mlx-c via CMake
+    let dst = cmake::Config::new(&mlx_c_dir)
+        .define("MLX_BUILD_TESTS", "OFF")
+        .define("MLX_BUILD_EXAMPLES", "OFF")
+        .define("MLX_BUILD_BENCHMARKS", "OFF")
+        .define("BUILD_SHARED_LIBS", "OFF")
+        .build();
+
+    // Link paths
+    let lib_dir = dst.join("lib");
+    println!("cargo:rustc-link-search=native={}", lib_dir.display());
+
+    // Also check lib64 (some CMake configs use this)
+    let lib64_dir = dst.join("lib64");
+    if lib64_dir.exists() {
+        println!("cargo:rustc-link-search=native={}", lib64_dir.display());
+    }
+
+    // Link mlx-c and mlx static libraries
+    println!("cargo:rustc-link-lib=static=mlxc");
+    println!("cargo:rustc-link-lib=static=mlx");
+
+    // Link macOS system frameworks required by MLX
+    println!("cargo:rustc-link-lib=framework=Metal");
+    println!("cargo:rustc-link-lib=framework=Foundation");
+    println!("cargo:rustc-link-lib=framework=Accelerate");
+    println!("cargo:rustc-link-lib=framework=MetalPerformanceShaders");
+
+    // Link C++ standard library
+    println!("cargo:rustc-link-lib=c++");
+
+    // Rerun if mlx-c sources change
+    println!("cargo:rerun-if-changed=mlx-c/CMakeLists.txt");
+    println!("cargo:rerun-if-changed=mlx-c/mlx/c/");
+}