up

Author: metascroy

CMakePresets.json

@@ -289,7 +289,8 @@
       "inherits": ["common"],
       "cacheVariables": {
         "EXECUTORCH_BUILD_PRESET_FILE": "${sourceDir}/tools/cmake/preset/mlx.cmake",
-        "CMAKE_OSX_DEPLOYMENT_TARGET": "14.0"
+        "CMAKE_OSX_DEPLOYMENT_TARGET": "14.0",
+        "CMAKE_CXX_FLAGS": "-DABSL_USES_STD_STRING_VIEW"
       },
       "condition": {
         "lhs": "${hostSystemName}",

backends/apple/mlx/examples/whisper/run_whisper.py

@@ -256,7 +256,7 @@ def run_whisper_inference(  # noqa: C901
     cache_position = cache_position + decoder_input_ids.shape[1]
 
     # Generation loop
-    for step in range(max_new_tokens):
+    for _step in range(max_new_tokens):
         current_pos = cache_position.item()
 
         # Check for forced token at this position

backends/apple/mlx/examples/whisper/run_whisper_hf.py

@@ -222,7 +222,7 @@ def run_inference(  # noqa: C901
     cache_position = cache_position + decoder_input_ids.shape[1]
 
     # Generation loop
-    for step in range(max_new_tokens):
+    for _step in range(max_new_tokens):
         current_pos = cache_position.item()
 
         # Check for forced token at this position

backends/apple/mlx/ops.py

@@ -1277,20 +1277,18 @@ def _unsqueeze_handler(P: MLXProgramBuilder, n: Node) -> Slot:
 @REGISTRY.register(
     target=[torch.ops.aten.squeeze.dims, torch.ops.aten.squeeze_copy.dims]
 )
-def _squeeze_handler(P: MLXProgramBuilder, n: Node) -> Slot:
+def _squeeze_dims_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     """Handle squeeze operation for specific dimensions.
 
     Removes dimensions of size 1 from the tensor at specified positions.
-    If dims is empty, removes all dimensions of size 1.
     """
     args = P.args(n)
     require_args(args, 2, 2, "aten.squeeze.dims")
     require_kwargs(P.kwargs(n), set(), "aten.squeeze.dims")
     x, dims = args
     out = P.make_or_get_slot(n)
 
-    # dims is typically a list of ints
-    dims_list = list(dims) if dims is not None else []
+    dims_list = list(dims) if dims is not None else None
 
     P.emit(
         SqueezeNode(
@@ -1302,6 +1300,30 @@ def _squeeze_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     return out
 
 
+@REGISTRY.register(
+    target=[torch.ops.aten.squeeze.default, torch.ops.aten.squeeze_copy.default]
+)
+def _squeeze_default_handler(P: MLXProgramBuilder, n: Node) -> Slot:
+    """Handle squeeze operation without specified dimensions.
+
+    Removes all dimensions of size 1 from the tensor.
+    """
+    args = P.args(n)
+    require_args(args, 1, 1, "aten.squeeze.default")
+    require_kwargs(P.kwargs(n), set(), "aten.squeeze.default")
+    (x,) = args
+    out = P.make_or_get_slot(n)
+
+    P.emit(
+        SqueezeNode(
+            x=P.slot_to_tid(x),
+            out=P.slot_to_tid(out),
+            dims=None,
+        )
+    )
+    return out
+
+
 @REGISTRY.register(target=[torch.ops.aten.cat.default])
 def _cat_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     """Handle concatenation of a list of tensors.

backends/apple/mlx/program_builder.py

@@ -93,6 +93,7 @@ def get_aten_target(target):
     torch.ops.aten.unsqueeze_copy.default: torch.ops.aten.unsqueeze.default,
     torch.ops.aten.squeeze_copy.dim: torch.ops.aten.squeeze.dim,
     torch.ops.aten.squeeze_copy.dims: torch.ops.aten.squeeze.dims,
+    torch.ops.aten.squeeze_copy.default: torch.ops.aten.squeeze.default,
     torch.ops.aten.expand_copy.default: torch.ops.aten.expand.default,
     torch.ops.aten.alias_copy.default: torch.ops.aten.alias.default,
 }

backends/apple/mlx/runtime/MLXInterpreter.h

@@ -464,15 +464,24 @@ inline void exec_tanh(const TanhNode& n, ExecutionState& st, StreamOrDevice s) {
 inline void
 exec_squeeze(const SqueezeNode& n, ExecutionState& st, StreamOrDevice s) {
   const auto& x = st.const_tensor_ref(n.x);
-  auto dims_fb = n.dims;
+  const auto& dims_fb = n.dims;
 
   if (dims_fb.size() > 0) {
-    // Squeeze specific dimensions
+    // Squeeze specific dimensions, filtering out non-size-1 dims to match
+    // PyTorch semantics where squeeze on a non-size-1 dim is a no-op.
     std::vector<int> dims;
     for (auto d : dims_fb) {
-      dims.push_back(d);
+      int axis = d < 0 ? d + static_cast<int>(x.ndim()) : d;
+      if (axis >= 0 && axis < static_cast<int>(x.ndim()) &&
+          x.shape(axis) == 1) {
+        dims.push_back(d);
+      }
+    }
+    if (dims.size() > 0) {
+      st.set_tensor(n.out, squeeze(x, dims, s));
+    } else {
+      st.set_tensor(n.out, x);
     }
-    st.set_tensor(n.out, squeeze(x, dims, s));
   } else {
     // Squeeze all dimensions of size 1
     st.set_tensor(n.out, squeeze(x, s));

backends/apple/mlx/test/test_ops.py

@@ -475,12 +475,12 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]:
 class SqueezeModel(nn.Module):
     """Model that squeezes a tensor at specified dimensions."""
 
-    def __init__(self, dims: Tuple[int, ...]):
+    def __init__(self, dims: Optional[Tuple[int, ...]] = None):
         super().__init__()
         self.dims = dims
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if len(self.dims) == 0:
+        if self.dims is None:
             return torch.squeeze(x)
         else:
             return torch.squeeze(x, dim=self.dims)
@@ -495,12 +495,19 @@ class SqueezeTest(OpTestCase):
     atol = 1e-5
 
     def __init__(
-        self, shape: Tuple[int, ...] = (1, 3, 1, 4), dims: Tuple[int, ...] = (0, 2)
+        self,
+        shape: Tuple[int, ...] = (1, 3, 1, 4),
+        dims: Optional[Tuple[int, ...]] = (0, 2),
     ):
         self.shape = shape
         self.dims = dims
         shape_str = "x".join(str(s) for s in shape)
-        dims_str = "_".join(str(d) for d in dims) if dims else "all"
+        if dims is None:
+            dims_str = "all"
+        elif len(dims) == 0:
+            dims_str = "empty"
+        else:
+            dims_str = "_".join(str(d) for d in dims)
         self.name = f"squeeze_{shape_str}_dims{dims_str}"
 
     @classmethod
@@ -511,6 +518,10 @@ def get_test_configs(cls) -> List["SqueezeTest"]:
             cls(shape=(3, 1, 4), dims=(1,)),
             cls(shape=(1, 1, 8), dims=(0, 1)),
             cls(shape=(2, 1, 3, 1), dims=(1, 3)),
+            # Squeeze all singleton dims (no dims specified)
+            cls(shape=(1, 3, 1, 4), dims=None),
+            # Dims include non-size-1 axes (should be no-op for those axes)
+            cls(shape=(1, 1, 1, 8198), dims=(0, 1, 2, 3)),
         ]
 
     def create_inputs(self) -> Tuple[torch.Tensor, ...]:

examples/models/parakeet/README.md

@@ -27,7 +27,6 @@ python export_parakeet_tdt.py --audio /path/to/audio.wav
 | `--output-dir` | Output directory for exports (default: `./parakeet_tdt_exports`) |
 | `--backend` | Backend for acceleration: `portable`, `xnnpack`, `metal`, `mlx`, `cuda`, `cuda-windows` (default: `xnnpack`) |
 | `--dtype` | Data type: `fp32`, `bf16`, `fp16` (default: `fp32`). Metal backend supports `fp32` and `bf16` only (no `fp16`). |
-| `--quantize` | Quantization mode: `int4` for int4 weight-only quantization via TorchAO (default: none) |
 | `--audio` | Path to audio file for transcription test |
 
 **Note:** The preprocessor is always lowered with the portable backend regardless of the `--backend` setting.
@@ -135,18 +134,20 @@ This generates:
 
 ### MLX Export (macOS)
 
-Export with MLX backend:
+Export with MLX backend (bf16, int4 quantized, group size 128):
 ```bash
-python export_parakeet_tdt.py --backend mlx --output-dir ./parakeet_mlx
-```
-
-Export with int4 quantization (reduces model size ~4x):
-```bash
-python export_parakeet_tdt.py --backend mlx --quantize int4 --output-dir ./parakeet_mlx_int4
+python export_parakeet_tdt.py \
+    --backend mlx \
+    --dtype bf16 \
+    --qlinear_encoder 4w \
+    --qlinear_encoder_group_size 128 \
+    --qlinear 4w \
+    --qlinear_group_size 128 \
+    --output-dir ./parakeet_mlx_4w
 ```
 
 This generates:
-- `parakeet_tdt.pte` - The compiled model with MLX delegate
+- `model.pte` - The compiled model with MLX delegate (~470 MB)
 - `tokenizer.model` - SentencePiece tokenizer
 
 ## C++ Runner
@@ -172,23 +173,18 @@ Then build the parakeet runner:
 cd examples/models/parakeet
 
 # CPU/XNNPACK build
-make parakeet-cpu
+cmake --workflow --preset parakeet-cpu
 
 # Metal build (macOS)
-make parakeet-metal
+cmake --workflow --preset parakeet-metal
 
 # CUDA build (Linux)
-make parakeet-cuda
+cmake --workflow --preset parakeet-cuda
 
 # MLX build (macOS)
-make parakeet-mlx
+cmake --workflow --preset parakeet-mlx
 ```
 
-Available presets:
-- `parakeet-cpu` - CPU-only build
-- `parakeet-cuda` - CUDA acceleration (Linux/Windows)
-- `parakeet-metal` - Metal acceleration (macOS)
-- `parakeet-mlx` - MLX acceleration (macOS)
 ### Running
 
 From the executorch root directory:
@@ -212,6 +208,12 @@ DYLD_LIBRARY_PATH=/usr/lib ./cmake-out/examples/models/parakeet/parakeet_runner
   --data_path examples/models/parakeet/parakeet_cuda/aoti_cuda_blob.ptd \
   --audio_path /path/to/audio.wav \
   --tokenizer_path examples/models/parakeet/parakeet_cuda/tokenizer.model
+
+# MLX
+./cmake-out/examples/models/parakeet/parakeet_runner \
+  --model_path examples/models/parakeet/parakeet_mlx_4w/model.pte \
+  --audio_path /path/to/audio.wav \
+  --tokenizer_path examples/models/parakeet/parakeet_mlx_4w/tokenizer.model
 ```
 
 ### Runner Arguments

examples/models/parakeet/export_parakeet_tdt.py

@@ -1,7 +1,6 @@
 """Export nvidia/parakeet-tdt-0.6b-v3 components to ExecuTorch."""
 
 import argparse
-import logging
 import os
 import shutil
 import tarfile
@@ -20,8 +19,6 @@
 from executorch.exir.passes import MemoryPlanningPass
 from torch.export import Dim, export
 
-logger = logging.getLogger(__name__)
-
 
 def load_audio(audio_path: str, sample_rate: int = 16000) -> torch.Tensor:
     """Load audio file and resample to target sample rate."""
@@ -442,7 +439,6 @@ def export_all(
         strict=False,
     )
 
-
     sample_rate = model.preprocessor._cfg.sample_rate
     window_stride = float(model.preprocessor._cfg.window_stride)
     encoder_subsampling_factor = int(getattr(model.encoder, "subsampling_factor", 8))
@@ -564,20 +560,13 @@ def _create_cuda_partitioners(programs, is_windows=False):
 
 
 def _create_mlx_partitioners(programs):
-    """Create MLX partitioners for all programs except preprocessor."""
+    """Create MLX partitioners for all programs."""
     from executorch.backends.apple.mlx.partitioner import MLXPartitioner
 
     print("\nLowering to ExecuTorch with MLX...")
 
     partitioner = {}
     for key in programs.keys():
-        # if key == "preprocessor":
-        #     # Skip preprocessor - FFT ops are not supported by MLX and fall back
-        #     # to portable pocketfft implementation. There is a bug in pocketfft
-        #     # that causes SIGABRT ("pointer being freed was not allocated") in
-        #     # release builds but not debug builds.
-        #     partitioner[key] = []
-        # else:
         partitioner[key] = [MLXPartitioner()]
 
     return partitioner, programs
@@ -621,38 +610,6 @@ def lower_to_executorch(programs, metadata=None, backend="portable"):
     )
 
 
-def apply_quantization(model, quantize: str) -> None:
-    """Apply quantization to the model using TorchAO.
-
-    Args:
-        model: The model to quantize
-        quantize: Quantization method ("int4" or "int8")
-    """
-    try:
-        from torchao.quantization.granularity import PerGroup
-        from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_
-    except ImportError:
-        logger.error("TorchAO not installed. Run: pip install torchao")
-        raise
-
-    logger.info(f"Applying {quantize} quantization to linear layers...")
-
-    if quantize == "int4":
-        quantize_(
-            model,
-            IntxWeightOnlyConfig(weight_dtype=torch.int4, granularity=PerGroup(128)),
-            lambda m, fqn: isinstance(m, torch.nn.Linear),
-        )
-    elif quantize == "int8":
-        quantize_(
-            model,
-            IntxWeightOnlyConfig(weight_dtype=torch.int8, granularity=PerGroup(128)),
-            lambda m, fqn: isinstance(m, torch.nn.Linear),
-        )
-    else:
-        logger.warning(f"Unknown quantization method: {quantize}")
-
-
 def main():
 
     parser = argparse.ArgumentParser()
@@ -729,13 +686,6 @@ def main():
         help="Group size for embedding quantization (default: 0 = per-axis)",
     )
 
-    parser.add_argument(
-        "--quantize",
-        type=str,
-        choices=["int4", "int8"],
-        default=None,
-        help="Quantization method for linear layers (requires torchao)",
-    )
     args = parser.parse_args()
 
     # Validate dtype
@@ -764,10 +714,6 @@ def main():
         print("Converting model to float16...")
         model = model.to(torch.float16)
 
-    # Apply quantization if requested
-    if args.quantize:
-        apply_quantization(model, args.quantize)
-
     print("\nExporting components...")
     export_dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float
     programs, metadata = export_all(

examples/models/parakeet/quantize.py

@@ -109,6 +109,7 @@ def linear_filter(m, fqn):
                 config = IntxWeightOnlyConfig(
                     weight_dtype=torch.int4,
                     granularity=granularity,
+                    intx_choose_qparams_algorithm="hqq_scale_only",
                 )
         elif qlinear_config == "8w":
             config = IntxWeightOnlyConfig(