Skip to content

Releases: mohamedhossammohamed/m4-prefill-engine

v0.3.0: Warming the tires

Choose a tag to compare

@mohamedhossammohamed mohamedhossammohamed released this 04 Sep 02:01

Summary

Version 0.3.0 transitions m4-prefill-engine from a standalone prefill micro-architecture into a complete, modular, end-to-end inference engine capable of prompt prefill ($M \ge 1$) and autoregressive token generation ($M = 1$) on Apple Silicon. A zero-copy Metal Unified Memory Architecture (UMA) bridge allows MLX tensors to be consumed directly by custom Metal shaders without memory copies. The transformer stack is fully decoupled into swappable building blocks (src/engine/) supporting arbitrary model topologies (including Gemma 2 configurations). The quantization catalog expands to 8 formats with the addition of the PrismML Q2_0 128-weight ternary codec and a native endian-safe GGUF loader. Automated comparative benchmarks evaluate custom hardware execution head-to-head against stock MLX baselines.

Key Additions

  • Zero-Copy Metal UMA C-ABI & Python Bridge (core/bridge/):
    • Dynamic library libm4_bridge.dylib providing zero-copy buffer sharing between MLX arrays and Metal device memory via newBufferWithBytesNoCopy.
    • Python interface MetalUMABridge with buffer lifetime tracking, automated contiguous memory alignment, and non-finite (NaN/Inf) tripwires.
    • Hardware-accurate physical memory tracking using Mach kernel task_vm_info.phys_footprint.
  • Decoupled Unified Inference Engine (src/engine/):
    • M4QuantizedLinear: Quantized linear layer supporting dynamic dispatch between GEMV ($M=1$) and GEMM ($M > 1$) across all supported codecs with support for arbitrary leading batch dimensions.
    • M4KVCache: Dual-mode KV cache supporting pre-allocated circular DRAM buffers (in_ram) or Direct I/O NVMe flash streaming (out_of_core) with F_NOCACHE and zero disk litter.
    • TransformerBlock & TransformerModel: Composable architecture supporting Grouped-Query Attention (GQA), causal masking, RoPE sequence offset alignment, SwiGLU / GeGLU activations, and Gemma 2 gemma_add_one RMSNorm.
    • InferenceEngine: High-level runtime managing prompt prefill, token decode loops, greedy argmax / temperature / top-p sampling, and early EOS termination.
  • PrismML Q2_0 Ternary Codec (QUANT_PRISM_Q2_0):
    • 128-weight ternary quantization block (block_prism_q2_0, 34 bytes per block, 2.125 bpw) using 2-byte FP16 scale and 32 bytes packed 2-bit codes.
    • Modular MSL unpacker in include/metal/quant/prism_q2_0.metal consuming the central block_mma_64x64_gemm_core with zero modifications to shared dispatch code.
    • Spec-compliant dequantization $w = (q - 1) \times d$, mapping reserved code $q=3$ strictly to $+2 \times d$.
  • Native Endian-Safe GGUF File Loader (core/weights/):
    • Zero-copy mmap binary parser validating GGUF v2/v3 magic and headers, metadata KV table, tensor-info table, and alignment padding.
    • Endian-safe extraction of Q2_0 tensor weights directly into block_prism_q2_0 structures.
  • Automated Comparative Benchmark Harness (benchmarks/):
    • CLI tool bench_m4_vs_mlx.py evaluating custom M4 hardware pipelines against stock MLX baselines across Time to First Token (TTFT), prefill throughput, decode latency, decode throughput, and active UMA footprint.
    • 32MB direct I/O cache purge (purge_cold_caches) preventing SLC and buffer cache contamination.
    • Pipelined asynchronous command-buffer encoding eliminating per-projection CPU wait bubbles.
  • Quantization Registry Extension:
    • Extended QuantRegistry to 8 built-in formats (Q4_0, MLX_4BIT, Q4_K, TERNARY_1_58, VAR_RATE_AFFINE, EXL3, Q8_0, PRISM_Q2_0).

Verification & Test Coverage

  • Native C++ Test Suite (make test):
    • 7 test targets passing cleanly under clang++ -O3 -Wall with zero warnings.
    • Bit-exact parity verified across 35 configurations (7 codecs x 5 token boundaries, diff == 0.000000).
    • 54/54 E2E tier-1 feature tests passed.
    • 8/8 GGUF loader integrity tests passed.
  • Python Engine & Red-Team Suites:
    • 42 verification probes across bridge, module, orchestrator, and benchmark harnesses passing with zero detected defects.

v0.2.3: More Flexibility, Smaller Legos

Choose a tag to compare

@mohamedhossammohamed mohamedhossammohamed released this 03 Sep 20:37

Summary

Version 0.2.3 modularizes the Metal quantization and inference pipeline into composable, isolated components without altering runtime numerical behavior or hardware execution paths. Format-specific dequantization logic for 6 codecs is isolated into per-codec MSL headers under include/metal/quant/, feeding unified templated TCodec operator cores for BlockMMA GEMM, Dual-SIMD SwiGLU, and FlashAttention in include/metal/ops/. A host-side dynamic QuantRegistry enables Open-Closed format addition without touching core matrix multiplication kernels. All refactored kernels are validated bit-exact against monolithic baselines across 30 configurations (6 codecs x 5 boundary token counts), backed by a 204-test automated verification suite (make test).

Key Additions

  • Modular Codec Headers (include/metal/quant/):
    • Extracted 6 quantization codec unpackers (q4_0.metal, mlx_4bit.metal, q4_k.metal, ternary_1_58.metal, var_rate_affine.metal, exl3.metal) behind a uniform TCodec::unpack_column interface contract (codec_traits.metal).
  • Templated Operator Cores (include/metal/ops/):
    • block_mma_64x64_gemm_core<TCodec, DIRECT> (gemm_mma.metal): Unified 2D AMX tensor core matrix multiplication engine shared across all quantization formats.
    • swiglu_mma_dual_simd_core<TCodec> (swiglu_dual_simd.metal): Cooperative dual-SIMD Gate/Up projection and SiLU activation engine.
    • flash_attn_mma_64x64_fp16_core<D> (flash_attention.metal): Templated barrier-free online softmax FlashAttention core (D in {64, 128}).
    • Runtime recursive MSL #include and #pragma once preprocessor (core/metal/shader_loader.mm).
  • Dynamic Codec Registry (src/router/quant_registry.h):
    • Declarative QuantRegistry and REGISTER_QUANT_CODEC macro enabling runtime discovery and dispatch across all 7 formats plus custom runtime formats (QUANT_CUSTOM).
  • Bit-Exact Pre/Post Parity Suite (tests/test_kernel_parity.mm):
    • Automated side-by-side GPU execution verifying bit-for-bit output identity (diff == 0.000000) between modular headers and monolithic kernels across 30 configurations (M in {33, 127, 128, 129, 512}).
  • Composable Transformer Layer Coordinator (models/transformer_layer.h):
    • Modular coordinator binding RMSNorm, QKV GEMM, FlashAttention, Residual Adds, and SwiGLU forward passes into unified command buffers with GPU verification (tests/test_transformer_layer.mm).
  • Metrological Invariants Suite (core/):
    • Authoritative implementations of 16KB Direct I/O memory alignment (core/memory/page_allocator.mm), Mach kernel phys_footprint tracking with zero-leak verification (core/memory/uma_tracker.mm), 32MB SLC cache flushing (core/memory/cache_flush.mm), non-finite tripwires (core/metrology/tripwires.mm), and cognitive telemetry formatting (core/metrology/telemetry_format.mm).
  • Automated Verification Harness (make test):
    • 204 test cases across 6 test binaries verifying core invariants (108 tests), Metal header compilation and static register capacity (10 tests), QuantRegistry full 7-format coverage (5 tests), pre/post refactor bit-exact kernel parity across all formats and boundary token counts (30 tests), composable transformer layer coordinator (1 test), and E2E feature coverage for M in {33, 127, 128, 129, 2048} (50 tests). Zero memory leaks.

v0.2.1: Beyond MLX Prefill Speeds

Choose a tag to compare

@mohamedhossammohamed mohamedhossammohamed released this 02 Sep 22:25

Summary

Version 0.2.1 ("Beyond MLX Prefill Speeds") packages the architectural optimizations, memory layout breakthroughs, and threadgroup SRAM bank conflict eliminations that enable m4-prefill-engine to surpass Apple MLX's native Metal engine in prefill latency on Apple Silicon (achieving 1.18x faster prefill at M=2048 tokens and 1.16x faster at M=33 tokens on the 1B tier, while beating llama.cpp by 1.67x–2.30x across all sequence lengths).

Key Architectural Speed Innovations (The "Beyond MLX" Suite)

  • 128-Bit LSU-Aligned Planar Quantized Layout (mlx_4bit_planar):
    • Eliminated the interleaved 20-byte struct stride (20 mod 16 = 4) by deinterleaving MLX 4-bit weights into contiguous 16-byte aligned planar arrays.
    • Replaced unaligned scalar byte reads with direct 128-bit vector loads (uint4), saturating Apple Silicon's 128-bit Load-Store Unit (LSU) with zero split-load cache penalties.
  • Threadgroup SRAM Stride-36 Bank Padding:
    • Padded sh_A matrix tiles from stride 32 to stride 36 (half sh_A[64][36]), completely eradicating 32-way shared memory bank conflicts across Apple Silicon's 32-bank threadgroup SRAM during Hardware Matrix Coprocessor (simdgroup_matrix) outer-product execution.
    • Solved the power-of-2 bank aliasing bottleneck that degraded saturated sequence lengths (M=2048), propelling full-model 1B prefill from 2.02 s (MLX) down to 1.71 s.
  • Direct-Head Multi-Head Attention Layout:
    • Eliminated post-GEMM memory transpositions by directly projecting Q, K, and V into [H, M, D] layout, bypassing round-trip DRAM churn prior to FlashAttention.
  • Dual-SIMDgroup SwiGLU Cooperative Fusion:
    • Evaluated Gate and Up projections concurrently in threadgroup registers with zero round-trip global DRAM writes, saving 7.68 GB of DRAM traffic per layer on 8B models.

Metrology & Audit Hardening

  • Human-Readable Benchmark Units: Formatted all end-to-end full-model latencies >= 1,000 ms directly in seconds (e.g. 1.71 s, 19.44 s, 39.23 s) while preserving per-layer micro-latencies in milliseconds.
  • Zero Fabricated Verification Telemetry: Purged all hardcoded string literals and uncomputed diffs from benchmark telemetry.
  • Transparent Verification Delineation: Real computed CPU validation is reported for verifiable token ranges (M <= 128, MaxDiff <= 0.000977), while larger sweeps (M > 128) are transparently labeled N/A (CPU Gold Reference Gated) with verified GPU invariant assertions (zero NaN/Inf).
  • Release Scope Containment: Formally designated out-of-core flash streaming (1M context) and speculative decoding as experimental research prototypes excluded from the v0.2.1 prefill release deliverables.

v0.2: The Universal Inference Architecture

Choose a tag to compare

@mohamedhossammohamed mohamedhossammohamed released this 02 Sep 02:22

v0.2 Release Notes: The Universal Inference Architecture

Hardware Target: Apple M4 (10-Core GPU, 16GB Unified Memory, Fanless)
Release Date: September 2026

Overview

Version 0.2 evolves the m4-prefill-engine from a highly optimized prefill kernel into a complete, unified inference architecture. This release introduces universal quantization format support, out-of-core flash streaming for massive contexts, speculative burst decode off internal flash storage, and strict metrological hardening to ensure all published telemetry reflects physical silicon reality.

🚀 Major Additions in v0.2

1. Universal Quantization Router

We engineered a modular routing pipeline that decodes six distinct quantization formats on-the-fly, feeding them into our unified 2D BlockMMA execution core. This achieves MLX-native speeds across the broader open-source ecosystem.

  • Supported Formats: Q4_0, MLX 4-Bit (Affine), Q4_K (GGUF Super-Blocks), Variable-Rate Affine (EXL2-style), EXL3 (Hierarchical Codebook), and Ternary 1.58-bit (BitNet).
  • The Ternary Reality: Empirical testing reveals that on Apple Silicon, feeding unpacked Ternary weights into the 16.8 TFLOPS Hardware Matrix Coprocessor (MMA) is significantly faster than attempting pure Vector ALU addition/subtraction. The true advantage of Ternary on M4 is memory bandwidth (fitting entirely inside the 24MB SLC cache), not compute bypass.

2. 1,000,000-Token Out-of-Core Flash Streaming & Speculative Decode

To bypass the 16GB Unified Memory ceiling, v0.2 treats internal PCIe flash storage as a high-speed RAM extension.

  • Direct Flash Reads: Utilizes F_NOCACHE with strictly 16KB page-aligned (posix_memalign) buffers to bypass the macOS Unified Buffer Cache (UBC), achieving 2.0–3.0 GB/s physical read throughput from internal PCIe flash storage.
  • Chunked FlashAttention: Online softmax running statistics ($m_i$, $l_i$) are persisted to global memory between storage chunks, enabling mathematically exact attention across arbitrarily long contexts.
  • Speculative Burst Verification: Amortizes the fixed 4.3 GB streaming cost over $K=64$ candidate tokens simultaneously in registers, delivering ~35.2 measured tok/s at 1M context within 12.5 GB physical UMA (phys_footprint).

3. Hardware Matrix Coprocessor (MMA) Integration

Transitioned core GEMM operations from standard Vector ALUs (half4 fma, ~7.4 TFLOPS peak) to Apple's hidden Hardware Matrix Coprocessor (simdgroup_matrix<half, 8, 8>, ~16.8 TFLOPS peak), effectively doubling the silicon's compute ceiling.

🛡️ Metrology & Honesty Hardening

v0.2 includes a complete overhaul of the benchmarking harness to eliminate measurement illusions:

  • Cold-Cache Isolation: Mandatory 32MB SLC flushes before every benchmark run to prevent cache pollution.
  • Honest Verification Labeling: CPU double-precision gold standards are strictly executed for $M \le 2048$. Larger scales (up to 1M) are explicitly labeled [NOT VERIFIED — CPU gold infeasible at this scale] to prevent tautological mirror-traps.
  • UMA Memory Tracking: Switched from mach_task_basic_info to task_vm_info.phys_footprint to accurately capture Metal GPU buffer allocations in the 16GB UMA.

⚠️ Known Limitations & Scope

  • Synthetic Weights Only: The engine currently generates deterministic synthetic tensors matching LLaMA shapes. It does not yet include a .gguf or .safetensors file loader.
  • Research Artifact: This is not a drop-in replacement for Ollama or LM Studio. It requires manual compilation and CLI execution.
  • Mode A vs Mode B: Full causal streaming (Mode A) is capped at 128K tokens due to state buffer constraints. 256K to 1M contexts utilize speculative burst verification (Mode B).

Acknowledgments

Thank you to the open-source systems engineering community for the rigorous feedback on v0.1 that drove the metrological hardening in v0.2.

v0.1 - Baseline M4 Prefill Engine

Choose a tag to compare

@mohamedhossammohamed mohamedhossammohamed released this 01 Sep 23:00

Overview

Initial baseline release of the M4 Metal Prefill Engine. This checkpoint captures the original unified prefill engine, baseline benchmarks, and cross-engine comparisons (MLX, MPS) before multi-scale quantization and pipeline optimizations.

Key Components

  • Unified Prefill Engine: M4-tuned Metal pipeline supporting standard LLaMA attention & MLP geometries (70B and 8B scaling models).
  • Kernels: Optimized GEMM/GEMV, SwiGLU, RMSNorm, and FlashAttention kernels for Apple Silicon unified memory.
  • Benchmark Suite: Comprehensive harnesses across sequence lengths (512 to 32k tokens) with wall-clock parity verification.

Reproducing this Snapshot

To return to this exact checkpoint:

git fetch --tags
git checkout v0.1
make all