Repository navigation
Releases: mohamedhossammohamed/m4-prefill-engine
Release list
v0.3.0: Warming the tires
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 (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.dylibproviding zero-copy buffer sharing between MLX arrays and Metal device memory vianewBufferWithBytesNoCopy. - Python interface
MetalUMABridgewith 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.
- Dynamic library
-
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) withF_NOCACHEand zero disk litter. -
TransformerBlock&TransformerModel: Composable architecture supporting Grouped-Query Attention (GQA), causal masking, RoPE sequence offset alignment, SwiGLU / GeGLU activations, and Gemma 2gemma_add_oneRMSNorm. -
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.metalconsuming the centralblock_mma_64x64_gemm_corewith zero modifications to shared dispatch code. - Spec-compliant dequantization
$w = (q - 1) \times d$ , mapping reserved code$q=3$ strictly to$+2 \times d$ .
- 128-weight ternary quantization block (
-
Native Endian-Safe GGUF File Loader (
core/weights/):- Zero-copy
mmapbinary 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_0structures.
- Zero-copy
-
Automated Comparative Benchmark Harness (
benchmarks/):- CLI tool
bench_m4_vs_mlx.pyevaluating 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.
- CLI tool
-
Quantization Registry Extension:
- Extended
QuantRegistryto 8 built-in formats (Q4_0,MLX_4BIT,Q4_K,TERNARY_1_58,VAR_RATE_AFFINE,EXL3,Q8_0,PRISM_Q2_0).
- Extended
Verification & Test Coverage
- Native C++ Test Suite (
make test):- 7 test targets passing cleanly under
clang++ -O3 -Wallwith 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.
- 7 test targets passing cleanly under
- 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
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 uniformTCodec::unpack_columninterface contract (codec_traits.metal).
- Extracted 6 quantization codec unpackers (
- 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
#includeand#pragma oncepreprocessor (core/metal/shader_loader.mm).
- Dynamic Codec Registry (
src/router/quant_registry.h):- Declarative
QuantRegistryandREGISTER_QUANT_CODECmacro enabling runtime discovery and dispatch across all 7 formats plus custom runtime formats (QUANT_CUSTOM).
- Declarative
- 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}).
- Automated side-by-side GPU execution verifying bit-for-bit output identity (
- 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).
- Modular coordinator binding RMSNorm, QKV GEMM, FlashAttention, Residual Adds, and SwiGLU forward passes into unified command buffers with GPU verification (
- Metrological Invariants Suite (
core/):- Authoritative implementations of 16KB Direct I/O memory alignment (
core/memory/page_allocator.mm), Mach kernelphys_footprinttracking 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).
- Authoritative implementations of 16KB Direct I/O memory alignment (
- 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
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_Amatrix 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.
- Padded
- 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
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_NOCACHEwith 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_infototask_vm_info.phys_footprintto 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
.ggufor.safetensorsfile 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
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