The AirLLMLlamaMlx backend runs Llama-family models natively on Apple Silicon Macs using MLX instead of PyTorch, handling checkpoint sharding, RoPE tuning, and greedy/sampled decoding entirely in MLX. Setup requires both mlx and torch installed on Apple Silicon hardware; the backend auto-splits model checkpoints on first use and tracks RAM consumption across forward passes. MLX is Apple's open-source array framework optimized for Apple Silicon's unified memory architecture, enabling tensor operations to run on the GPU or Neural Engine without copying data across separate memory pools.
AirLLMLlamaMlx, defined in air_llm/airllm/airllm_llama_mlx.py, is the macOS/MLX backend for Llama-family models, importing mlx.core and mlx.nn in place of PyTorch for all forward computation.[1] macOS support requires both mlx and torch to be installed, and only Apple Silicon hardware is supported — Intel Macs are not.[2]
AirLLMLlamaMlx.__init__ calls find_or_create_local_splitted_path to split the checkpoint into per-layer shards on first use, sharing the same on-disk splitting step as the CUDA backend — see On-disk splitting and persistence.[1] AirLLMLlamaMlx defines a fixed layer-name mapping via set_layer_names_dict(): embed → model.embed_tokens, layer_prefix → model.layers, norm → model.norm, and lm_head → lm_head; subclasses can override this mapping to support other architectures.[1]
sanitize_config in air_llm/airllm/airllm_llama_mlx.py defaults rope_theta to 10000 when the key is absent from the config dict, matching the RoPE base used by the original LLaMA.[1] sanitize_config also defaults n_kv_heads to the full n_heads count when the key is absent, enabling GQA-unaware configs to function with standard MHA semantics.[1] get_model_args_from_config hardcodes rope_traditional=False, so the MLX backend always uses the non-traditional (non-GPT-J) RoPE rotation order regardless of what the model config specifies.[1]
AirLLMLlamaMlx uses psutil.virtual_memory() — not GPU memory — to measure available RAM, tracking consumed and peak-consumed memory when show_memory_util=True.[1] The sample function uses greedy decoding (mx.argmax) when temperature=0 and categorical sampling otherwise; temperature=0 is the default for both generate and model_generate.[1]
Sources