Merge pull request #5 from second-state/add-mlx-backend

Add Apple MLX backend for macOS

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,26 @@ 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-latest
+            name: Linux x86_64 (CUDA)
+            backend: tch
+            asset-name: asr-linux-x86_64-cuda
+            libtorch-url: https://download.pytorch.org/libtorch/cu128/libtorch-cxx11-abi-shared-with-deps-2.7.1%2Bcu128.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 +41,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 +61,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 +79,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