mac-code

repository·main·Indexed 21 days ago

https://github.com/walter-grace/mac-code

A project for running large language models (LLMs) on Apple Silicon Macs, specifically designed for models that exceed available system RAM. It implements techniques such as 2-bit quantization, 'Flash Streaming' (streaming weights from SSD to GPU), and the mlx-expert-sniper package (v0.2.0) for MoE models. The repository also includes mac-tensor for distributed MoE inference across multiple Macs via HTTP and a Vision Agent pairing Gemma 4 with Falcon Perception.

Tokens
47.8K
Snippets
145
Records
227
Agent score
76%

What's inside mac-code

  1. Overview of Expert Sniper on d-Matrix Corsair

    main

    Expert Sniper is a Mixture of Experts (MoE) inference system designed to map onto d-Matrix Corsair hardware. The system optimizes MoE inference by streaming only active experts from slow storage (LPDDR5) into fast memory (SRAM), significantly reducing the hardware footprint required for large models. For example, a 122B MoE model that would typically require 64 Corsair cards using dense inference can be run on just 2 cards using Expert Sniper.

    Key optimization techniques include:

    • LRU expert cache: Maintains frequently used experts in fast memory.
    • Routing bias: Steers the router toward already-cached experts to achieve high hit rates (up to 97%).
    • Dead expert elimination: Skips experts that never fire (up to 45.2% of experts in some models).
    • Co-activation prefetch: Predicts next-layer experts before they are needed.
    • Union-of-experts batching: Deduplicates expert verification in batched tokens (up to 58% dedup).
  2. What is mac-tensor and how does it work?

    main

    mac-tensor is a system designed to run large Mixture-of-Experts (MoE) parameter models (like 35B models) across multiple Macs over a network using HTTP. It allows users to pool RAM from multiple machines without needing a GPU cluster or CUDA.

    Architecture

    • Coordinator: A single Mac that runs the attention, embeddings, and router weights (~1.5 GB RAM). It manages the orchestration.
    • Expert Nodes: Multiple Macs that hold partitions of the FFN (Feed-Forward Network) experts in RAM (~10 GB per node). They compute the FFN layers and return results to the coordinator via HTTP.

    When to use mac-tensor vs mlx-sniper

    • Use mac-tensor (Distributed) if you need to pool RAM across multiple Macs for models that are too large for a single machine (e.g., 70B+), or if you want all experts resident in RAM to avoid SSD latency.
    • Use mlx-sniper (Single Mac) if you have one powerful Mac. It is significantly faster (3-4x) because it streams experts from local NVMe using an LRU cache, avoiding network round-trip latencies.
  3. How Tiny Bit's vision tools work

    main

    Tiny Bit uses a tiered architecture for vision tasks to balance speed and accuracy. A fast 1-bit text model (like Bonsai) performs reasoning and discovers vision tools via search_tools. When an image is involved, it offloads to specialized backends:

    • vision_describe: Uses mac-tensor (Gemma 4-26B MoE) to return natural language descriptions of an image.
    • falcon_ground: Uses mac-tensor (Falcon Perception 0.6B) to return pixel-precise metadata like centroids, bounding boxes, and areas.

    This allows the agent to perform complex spatial reasoning (e.g., "Which bird is closest?") by combining fast text reasoning with targeted, slower vision calls.

  4. Architecture of Distributed Expert Sniper

    main

    The Distributed Expert Sniper uses a coordinator-worker pattern to partition Mixture-of-Experts (MoE) models across multiple nodes.

    Coordinator Node (e.g., M2-1):

    • Handles the core model logic: embedding, attention (RMSNorm, self_attn/linear_attn), residual connections, and the final lm_head.
    • Runs the Router, which identifies the top-8 expert IDs per layer.
    • Dispatches requests to worker nodes based on the expert IDs.
    • Receives expert FFN (Feed-Forward Network) results and performs the weighted sum.

    Worker Node (e.g., M2-2):

    • Holds a specific partition of experts in RAM (e.g., Experts 128-255).
    • Receives (layer_idx, expert_ids, hidden_state).
    • Computes gather_qmm FFN for the requested experts.
    • Returns the weighted expert output to the coordinator.

    This architecture allows running models larger than a single node's RAM by distributing expert partitions across a network.

  5. How MoE Expert Streaming works for low-RAM inference

    main

    MoE Expert Streaming allows running large models (like Qwen3.5-35B) on hardware with limited RAM (e.g., 8 GB) by exploiting model sparsity.

    Instead of loading the entire model into memory, the system:

    1. Pins essential weights (attention, routers, embeddings) in RAM (~1.4 GB).
    2. Streams only the active experts from an SSD on demand (~20.6 GB).

    Since only a small fraction of experts (e.g., 8 out of 256) are needed per token, the system can read only the necessary expert blocks from the SSD during the forward pass, overlapping I/O with computation to maintain usable token speeds.

  6. Understand the RAM budget for expert streaming

    main

    Expert streaming divides total system RAM into four primary consumers. To successfully run large MoE models on limited hardware, you must balance these components:

    1. macOS + apps: Typically requires ~2-3 GB of unavoidable overhead.
    2. Pinned weights: Fixed memory for attention, router, shared experts, and embeddings (e.g., ~1.4 GB for a 35B model).
    3. Compute buffers: Memory for KV cache, activations, and scratch space (~0.5-1 GB).
    4. Expert LRU cache: The remaining available RAM. Maximizing this is the most effective way to increase inference speed by improving the cache hit rate.
    5. Expert mmap pages: Managed by the OS; these are evicted when RAM is full.
  7. How Flash Streaming works

    main

    Flash Streaming allows you to run models that are larger than your available RAM by splitting the model into two parts based on access patterns:

    1. Pinned in RAM (4-6 GB): Contains Attention weights, embeddings, norms, and the KV cache. These are loaded once and stay in memory.
    2. Streamed from SSD per token: Contains the FFN (Feed-Forward Network) weights, which constitute the bulk of the model. These are loaded layer-by-layer from the SSD, used for a single matrix multiplication, and then discarded. This keeps memory usage flat.

    The process per token:

    • For each layer:
      1. Run attention (from RAM).
      2. Load FFN weights from SSD.
      3. Run FFN matmul on GPU.
      4. Discard FFN weights.

    MoE Optimization: For Mixture-of-Experts (MoE) models, only the active experts (e.g., 8 experts, ~14 MB) are loaded from SSD per layer, rather than the entire expert set, making MoE models significantly faster under Flash Streaming.

  8. Compare SRAM requirements for Dense vs MoE with Expert Sniper

    main

    Expert Sniper drastically reduces the amount of fast memory (SRAM) required for large MoE models compared to dense models. The following table illustrates the savings in terms of required cards (assuming expert SRAM is the primary driver):

    ModelDense SRAM NeededMoE + Sniper Active SetCards (Dense)Cards (Sniper)Savings
    35B17.5 GB1.41 GB919x
    122B61 GB3.51 GB31215x
    397B~200 GB7.0 GB100425x
    | Model | Dense SRAM Needed | MoE + Sniper Active Set | Cards (Dense) | Cards (Sniper) | Savings |
    |---|---|---|---|---|---|
    | 35B | 17.5 GB | 1.41 GB | 9 | **1** | **9x** |
    | 122B | 61 GB | 3.51 GB | 31 | **2** | **15x** |
    | 397B | ~200 GB | 7.0 GB | 100 | **4** | **25x** |
  9. How MLX Expert Sniper works

    main

    The project implements expert streaming to run MoE models larger than available RAM. The core mechanism involves:

    1. RAM Pinning: Keeping the attention mechanism, router, and shared experts resident in RAM (~0.87 GB).
    2. Expert Streaming: Streaming only the active experts (e.g., 8 out of 256) from the SSD using F_NOCACHE and pread.
    3. LRU Caching: An Least Recently Used (LRU) cache keeps 'hot' experts in memory, typically achieving an 85-88% hit rate.
    4. Kernel Fusion: Using gather_qmm to fuse quantized matrix multiplications across the active experts.
  10. Scale MoE inference via Thunderbolt Mac Clustering

    main

    To overcome the SSD bandwidth bottleneck when streaming experts for large MoE models, you can cluster multiple Macs using Thunderbolt to create a distributed expert pool. This uses Pipeline Parallelism, where each machine is responsible for a subset of the model's layers.

    Why it works

    • Low Interconnect Overhead: The data transferred between machines (activations) is tiny (~2 KB per token) compared to the expert data being loaded from local SSDs (~540 MB).
    • Linear Scaling: Because the bottleneck is local SSD bandwidth rather than the interconnect, adding more machines increases the total available bandwidth for the cluster.

    Scaling Example (using 16 GB Macs)

    MachinesTotal RAMModelLayers/MachineExpected tok/s
    1× 16 GB16 GBQ4_K_M (21 GB)401-5
    2× 16 GB32 GBQ4_K_M (21 GB)2010-20
    4× 16 GB64 GBQ4_K_M (21 GB)1020-30
  11. How LlamaCPPEngine and LlamaCPPOpenAIEngine work together

    main

    The module uses a two-tier adapter pattern to interface with local llama.cpp servers:

    1. LlamaCPPEngine (The Adapter): This is the high-level interface. It handles the logic of normalizing user input. It inspects the type of llm_input to decide whether to route to the /v1/completions or /v1/chat/completions endpoint. It also handles automatic model selection by querying the server for available models and picking the first one.

    2. LlamaCPPOpenAIEngine (The Concrete Engine): This is the low-level implementation. It uses an OpenAI client (configured to a local base_url) to perform the actual network calls. It handles the specific routing to /v1/models, /v1/chat/completions, or /v1/completions and manages the transformation of OpenAI response objects into the dictionary or SSE-style string formats expected by the worker.

  12. Gemma 4 Architecture details

    main

    Gemma 4 utilizes a hybrid attention architecture:

    • Sliding window attention: 1024 tokens on 5 of every 6 layers.
    • Full attention: Every 6th layer, utilizing K=V sharing.
    • Activation: gelu_pytorch_tanh.
    • Logit softcapping: Set at 30.0.

    The 26B-A4B MoE variant includes 128 experts with top-8 routing per layer, running in parallel with the dense MLP.