CoreML forced aligner — T=768, inline fp16 embedding gather, stride-aware MLMultiArray access, finite -1e4 causal mask

README.md

@@ -18,49 +18,71 @@ tags:
 - timestamps
 - neural-engine
 base_model: Qwen/Qwen3-ForcedAligner-0.6B
+base_model_relation: quantized
 pipeline_tag: audio-classification
 ---
 
 # Qwen3-ForcedAligner-0.6B — CoreML INT8
 
-CoreML conversion of [Qwen/Qwen3-ForcedAligner-0.6B](https://huggingface.co/Qwen/Qwen3-ForcedAligner-0.6B) with INT8 palettization for Apple Neural Engine.
+CoreML conversion of [Qwen/Qwen3-ForcedAligner-0.6B](https://huggingface.co/Qwen/Qwen3-ForcedAligner-0.6B) with INT8 kmeans palettization on the two transformer bundles, for Apple Silicon.
 
-Predicts word-level timestamps in a single forward pass.
+Predicts word-level timestamps for audio + text pairs in a single non-autoregressive forward pass.
 
-## Models
+## Files
 
-| Model | Description | Quantization |
-|-------|-------------|-------------|
-| `encoder.mlmodelc` | Audio encoder (24 layers) | INT8 palettized |
-| `decoder.mlmodelc` | Text decoder + classify head (28 layers) | INT8 palettized |
+| File | Role | Size |
+|---|---|---|
+| `audio_encoder.mlmodelc` / `.mlpackage` | 24-layer block-attention encoder (INT8-palettized); fixed 30 s mel input | ~304 MB |
+| `text_decoder.mlmodelc` / `.mlpackage` | 28-layer non-AR decoder + 5000-class classify head (INT8-palettized); fixed T=768 | ~427 MB |
+| `embed_tokens.fp16.bin` | Raw little-endian fp16 token embedding table, shape [152 064, 1024] (not palettized) | ~297 MB |
+| `config.json` | Runtime config — variant, classify_num, timestamp_segment_time, fixed shapes | small |
+| `vocab.json`, `merges.txt`, `tokenizer_config.json` | Qwen3 BPE tokenizer files | ~5 MB |
 
-## Usage
+The embed-tokens table is shipped as a raw fp16 binary rather than a CoreML package — the Swift runtime memory-maps the file and gathers rows with `vImageConvert_Planar16FtoPlanarF`, which costs ~0.5 ms per alignment instead of the ~70 ms an mlpackage round-trip would cost. Keeping the embedding at fp16 (not palettized) preserves classify-time accuracy across the full vocabulary; the palettization budget is spent on the two large transformer bundles, where it matters most for disk and ANE memory.
+
+## Usage (Swift via [speech-swift](https://github.com/soniqo/speech-swift))
 
 ```swift
+import Qwen3ASR
+
 let aligner = try await CoreMLForcedAligner.fromPretrained(
-    modelId: "aufklarer/Qwen3-ForcedAligner-0.6B-CoreML-INT8"
-)
-let aligned = aligner.align(audio: samples, text: "Hello world", sampleRate: 24000)
+    modelId: "aufklarer/Qwen3-ForcedAligner-0.6B-CoreML-INT8")
+let aligned = try aligner.align(
+    audio: pcmSamples,
+    text: "Can you guarantee that the replacement part will be shipped tomorrow?",
+    sampleRate: 16000,
+    language: "English")
 ```
 
-## Variants
+CLI: `speech align audio.wav --engine coreml --aligner-model aufklarer/Qwen3-ForcedAligner-0.6B-CoreML-INT8`
 
-| Variant | Backend | Size | Model ID |
-|---------|---------|------|----------|
-| CoreML INT4 | Neural Engine | ~630 MB | [aufklarer/Qwen3-ForcedAligner-0.6B-CoreML-INT4](https://huggingface.co/aufklarer/Qwen3-ForcedAligner-0.6B-CoreML-INT4) |
-| **CoreML INT8** | **Neural Engine** | **~1.0 GB** | **aufklarer/Qwen3-ForcedAligner-0.6B-CoreML-INT8** |
-| MLX 4-bit | GPU | ~979 MB | [aufklarer/Qwen3-ForcedAligner-0.6B-4bit](https://huggingface.co/aufklarer/Qwen3-ForcedAligner-0.6B-4bit) |
-| MLX 8-bit | GPU | ~1.4 GB | [aufklarer/Qwen3-ForcedAligner-0.6B-8bit](https://huggingface.co/aufklarer/Qwen3-ForcedAligner-0.6B-8bit) |
+## Performance (M2 Max, 64 GB, debug build, 20 s clip)
 
-## Links
+| Metric | Value |
+|---|---|
+| RTF | 0.014 (69× faster than real-time) |
+| Peak RSS | 697 MB |
+| Median alignment time | ~289 ms |
 
-- **Swift library**: [soniqo/speech-swift](https://github.com/soniqo/speech-swift)
-- **Base model**: [Qwen/Qwen3-ForcedAligner-0.6B](https://huggingface.co/Qwen/Qwen3-ForcedAligner-0.6B)
+Per-stage profile (`COREML_ALIGN_PROFILE=1`):
 
----
+```
+mel=150ms  encoder=78ms  embedding=0.5ms  splice=0ms  decoder=54ms  argmax=7ms
+```
 
----
+The decoder uses `MLComputeUnits = .all` because its 28 layers exceed the ~26-layer ANE graph cap; the encoder runs on `.cpuAndNeuralEngine`.
+
+## Architecture
+
+Same model as the upstream Qwen3-ForcedAligner. Inference is non-autoregressive: one forward pass over a chat-template prompt (system + audio + assistant + `<ts>` word `<ts>` slots), then argmax at the timestamp positions, then LIS monotonicity correction.
+
+The causal mask is baked into the exported graph as a constant with a finite ``-1e4`` fill value (not ``-inf``) so the fp16 softmax cannot produce NaN.
+
+Conversion source: [soniqo/speech-models/models/forced-aligner/export/convert_coreml.py](https://github.com/soniqo/speech-models).
+
+## Variants
+
+- This repo: **INT8** — kmeans-palettized text decoder + audio encoder, ~880 MB, similar RTF, ~35% lower peak memory than FP16.
+- Companion: [`aufklarer/Qwen3-ForcedAligner-0.6B-CoreML-FP16`](https://huggingface.co/aufklarer/Qwen3-ForcedAligner-0.6B-CoreML-FP16) — full precision, ~1.7 GB, highest fidelity.
 
-- **Guide**: [soniqo.audio/guides/align](https://soniqo.audio/guides/align)
-- **Docs**: [soniqo.audio](https://soniqo.audio)
-- **GitHub**: [soniqo/speech-swift](https://github.com/soniqo/speech-swift)
+For pure MLX paths (no CoreML), see the `4bit`, `5bit`, `8bit`, and `bf16` siblings under the same HuggingFace org.

audio_encoder.mlmodelc/analytics/coremldata.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c3cff651be40392d81b73d1c5a449a43a68eb7f5d4a642f0a2efaf315aa4d17e
+size 243

audio_encoder.mlmodelc/coremldata.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d310aae19a7b4bb18c5a9dbbeca9a4afc44062b3b28ef2df72bc3c273b39597e
+size 376

audio_encoder.mlmodelc/metadata.json

@@ -0,0 +1,102 @@
+[
+  {
+    "metadataOutputVersion" : "3.0",
+    "storagePrecision" : "Mixed (Float16, Palettized (8 bits))",
+    "outputSchema" : [
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 390 × 1024)",
+        "shortDescription" : "",
+        "shape" : "[1, 390, 1024]",
+        "name" : "audio_embeddings",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Int32",
+        "formattedType" : "MultiArray (Int32 1)",
+        "shortDescription" : "",
+        "shape" : "[1]",
+        "name" : "output_length",
+        "type" : "MultiArray"
+      }
+    ],
+    "modelParameters" : [
+
+    ],
+    "specificationVersion" : 8,
+    "mlProgramOperationTypeHistogram" : {
+      "Ios17.logicalAnd" : 2,
+      "Ios17.reshape" : 101,
+      "Ios16.softmax" : 24,
+      "Ios17.matmul" : 48,
+      "Ios17.transpose" : 98,
+      "Ios17.maximum" : 1,
+      "Select" : 1,
+      "Ios17.expandDims" : 4,
+      "Ios17.add" : 80,
+      "Ios17.gather" : 1,
+      "Ios17.layerNorm" : 49,
+      "Ios17.less" : 1,
+      "Ios17.sub" : 8,
+      "Ios16.constexprLutToDense" : 153,
+      "Ios17.conv" : 3,
+      "Ios17.clip" : 6,
+      "Ios17.linear" : 147,
+      "Ios16.gelu" : 28,
+      "Ios17.floorDiv" : 7,
+      "Ios17.minimum" : 1,
+      "Ios17.cast" : 11,
+      "Ios17.greater" : 6,
+      "Ios17.mul" : 26
+    },
+    "computePrecision" : "Mixed (Float16, Int32, UInt16)",
+    "isUpdatable" : "0",
+    "stateSchema" : [
+
+    ],
+    "availability" : {
+      "macOS" : "14.0",
+      "tvOS" : "17.0",
+      "visionOS" : "1.0",
+      "watchOS" : "10.0",
+      "iOS" : "17.0",
+      "macCatalyst" : "17.0"
+    },
+    "modelType" : {
+      "name" : "MLModelType_mlProgram"
+    },
+    "userDefinedMetadata" : {
+      "com.github.apple.coremltools.version" : "8.3.0",
+      "com.github.apple.coremltools.source" : "torch==2.3.1",
+      "com.github.apple.coremltools.source_dialect" : "TorchScript"
+    },
+    "inputSchema" : [
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float32",
+        "formattedType" : "MultiArray (Float32 1 × 128 × 3000)",
+        "shortDescription" : "",
+        "shape" : "[1, 128, 3000]",
+        "name" : "mel",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Int32",
+        "formattedType" : "MultiArray (Int32 1)",
+        "shortDescription" : "",
+        "shape" : "[1]",
+        "name" : "mel_length",
+        "type" : "MultiArray"
+      }
+    ],
+    "generatedClassName" : "audio_encoder",
+    "method" : "predict"
+  }
+]

audio_encoder.mlmodelc/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7b4eddd9b588c002922922840086d08a7b195a6932774c506dc2d4cdc94a1b9b
+size 317035136

audio_encoder.mlpackage/Data/com.apple.CoreML/model.mlmodel

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2cd2e45fc8711b8133e4ddfa80efc4ee8298255eba8af178b26b7c75f0d64e4b
-size 264149
+oid sha256:a8c09d520ef68809a3c3b313c484ab68b1787cc7f2cf9905106ffd60553b95cb
+size 437954

audio_encoder.mlpackage/Data/com.apple.CoreML/weights/weight.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a17493ebfda4a4cc21fc5928e03fa95286ddb65443689be38b8e06f07cd7ea51
-size 316866944
+oid sha256:7b4eddd9b588c002922922840086d08a7b195a6932774c506dc2d4cdc94a1b9b
+size 317035136

audio_encoder.mlpackage/Manifest.json

@@ -1,18 +1,18 @@
 {
     "fileFormatVersion": "1.0.0",
     "itemInfoEntries": {
-        "6D563B8A-0732-4B3B-A118-5582764B884B": {
-            "author": "com.apple.CoreML",
-            "description": "CoreML Model Weights",
-            "name": "weights",
-            "path": "com.apple.CoreML/weights"
-        },
-        "8D8287BF-F06B-4E15-97EF-5CCF1702CFF9": {
+        "8C63C029-A00C-40CD-A77E-0D338CBE599A": {
             "author": "com.apple.CoreML",
             "description": "CoreML Model Specification",
             "name": "model.mlmodel",
             "path": "com.apple.CoreML/model.mlmodel"
+        },
+        "A771C53E-C476-4479-A9D9-3413760251AC": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
         }
     },
-    "rootModelIdentifier": "8D8287BF-F06B-4E15-97EF-5CCF1702CFF9"
+    "rootModelIdentifier": "8C63C029-A00C-40CD-A77E-0D338CBE599A"
 }

config.json

@@ -1,7 +1,11 @@
 {
   "format": "coreml",
+  "variant": "int8",
   "palettization_bits": 8,
   "hidden_size": 1024,
   "classify_num": 5000,
-  "timestamp_segment_time": 0.08
+  "timestamp_segment_time": 0.08,
+  "fixed_mel_frames": 3000,
+  "fixed_audio_tokens": 390,
+  "text_fixed_t": 768
 }

embed_tokens.fp16.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a6968a6f3a6520111cf9288239a09e766cdebba9067c43e59ef06c6aac8e0512
+size 311427072

text_decoder.mlmodelc/analytics/coremldata.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ac45c534534d1b8c4c28c3823a65ca73e5c511d412bdd712feec35aae77d1c6e
+size 243

text_decoder.mlmodelc/coremldata.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:147c6ed553d4ebb7ff9f1cc358416051a08e92a90d94f63c63f1bf8aaa5d96e2
+size 323

text_decoder.mlmodelc/metadata.json

@@ -0,0 +1,76 @@
+[
+  {
+    "metadataOutputVersion" : "3.0",
+    "storagePrecision" : "Mixed (Float16, Palettized (8 bits))",
+    "outputSchema" : [
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 768 × 5000)",
+        "shortDescription" : "",
+        "shape" : "[1, 768, 5000]",
+        "name" : "logits",
+        "type" : "MultiArray"
+      }
+    ],
+    "modelParameters" : [
+
+    ],
+    "specificationVersion" : 8,
+    "mlProgramOperationTypeHistogram" : {
+      "Ios17.mul" : 506,
+      "Ios17.linear" : 197,
+      "Ios17.transpose" : 252,
+      "Ios17.sub" : 56,
+      "Ios16.constexprLutToDense" : 200,
+      "Ios17.rsqrt" : 113,
+      "Ios17.matmul" : 56,
+      "Ios17.concat" : 56,
+      "Ios17.sliceByIndex" : 112,
+      "Ios17.add" : 253,
+      "Ios16.reduceMean" : 113,
+      "Tile" : 56,
+      "Ios17.pow" : 113,
+      "Ios16.softmax" : 28,
+      "Ios17.reshape" : 224,
+      "Ios17.cast" : 1,
+      "Ios16.silu" : 28
+    },
+    "computePrecision" : "Mixed (Float16, Float32, Int32)",
+    "isUpdatable" : "0",
+    "stateSchema" : [
+
+    ],
+    "availability" : {
+      "macOS" : "14.0",
+      "tvOS" : "17.0",
+      "visionOS" : "1.0",
+      "watchOS" : "10.0",
+      "iOS" : "17.0",
+      "macCatalyst" : "17.0"
+    },
+    "modelType" : {
+      "name" : "MLModelType_mlProgram"
+    },
+    "userDefinedMetadata" : {
+      "com.github.apple.coremltools.version" : "8.3.0",
+      "com.github.apple.coremltools.source" : "torch==2.3.1",
+      "com.github.apple.coremltools.source_dialect" : "TorchScript"
+    },
+    "inputSchema" : [
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float32",
+        "formattedType" : "MultiArray (Float32 1 × 768 × 1024)",
+        "shortDescription" : "",
+        "shape" : "[1, 768, 1024]",
+        "name" : "inputs_embeds",
+        "type" : "MultiArray"
+      }
+    ],
+    "generatedClassName" : "text_decoder",
+    "method" : "predict"
+  }
+]

text_decoder.mlmodelc/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:083cfeabb070794b7745694564f0154aa4329020c7eb344aff3406f348b25e3f
+size 446498960

text_decoder.mlpackage/Data/com.apple.CoreML/model.mlmodel

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:086eca33cb9adb1a88ced5d058b1674ae7c772ae93cf3cb09309ea07ce0d5d37
-size 750916
+oid sha256:021f16e825a94b093779ca67abc42d5097e1fd75f56b45346b705602be610fc0
+size 667265

text_decoder.mlpackage/Data/com.apple.CoreML/weights/weight.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0b77e01591896e672e5e595b472b9afa63348fa8438fb35066e9702e97d3e521
-size 446334480
+oid sha256:083cfeabb070794b7745694564f0154aa4329020c7eb344aff3406f348b25e3f
+size 446498960

text_decoder.mlpackage/Manifest.json

@@ -1,18 +1,18 @@
 {
     "fileFormatVersion": "1.0.0",
     "itemInfoEntries": {
-        "50FFC863-A9F2-4984-BBB5-AEE058C2841C": {
-            "author": "com.apple.CoreML",
-            "description": "CoreML Model Specification",
-            "name": "model.mlmodel",
-            "path": "com.apple.CoreML/model.mlmodel"
-        },
-        "5E01B5E9-3ED2-4433-B70A-789523FA4835": {
+        "0C8546BF-7845-4333-8C21-8CE8474C52B6": {
             "author": "com.apple.CoreML",
             "description": "CoreML Model Weights",
             "name": "weights",
             "path": "com.apple.CoreML/weights"
+        },
+        "42F046A4-C822-4051-BD7B-791103EB03DD": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
         }
     },
-    "rootModelIdentifier": "50FFC863-A9F2-4984-BBB5-AEE058C2841C"
+    "rootModelIdentifier": "42F046A4-C822-4051-BD7B-791103EB03DD"
 }