TT-AMX, a zero-copy Tensor-Train inference engine for Apple Silicon

TT-AMX, a bare-metal C++ engine for Apple Silicon, runs Tensor-Train compressed linear layers on the undocumented AMX coprocessor, achieving a 2x speedup over dense FP16 GEMV under cold-cache conditions with half the memory. The engine, which uses zero-copy permutations and asymmetric factorization, reaches 889–947 GFLOP/s (62–66% of the AMX ceiling) and is open-source on GitHub.

A bare-metal C++ engine that runs Tensor-Train TT/MPO compressed linear layers on Apple Silicon by dispatching the contraction onto the undocumented AMX coprocessor . At 4x compression it beats a dense FP16 GEMV by ~2x under realistic cold-cache conditions , using half the memory. Tensor-Train is attractive for on-device work because it is a continuously differentiable compression manifold — unlike discrete INT4 quantization — which makes it usable for on-device PEFT, continual learning, and quantum-inspired ML. Historically it has been unusable for a different reason: a 5–10x latency penalty. That penalty turns out to be arithmetic, not cache. A TT matvec does 3–35x more multiply-accumulates than the dense matvec it replaces measured with numpy's optimal contraction order . No amount of cache tuning changes a MAC count. This engine wins anyway, by making those extra MACs nearly free. 1. Zero-copy ahead-of-time core permutation. A single offline transpose 1,2,0 on core 1 turns the runtime into two back-to-back cblas sgemm calls with nothing between them — no loop, no permute, no copy. The intermediate is reinterpreted, not moved. 2. Asymmetric factorization. Balanced tensor shapes are the wrong default. Sweeping 81 factorization pairs on real Qwen2.5-1.5B weights: | m factors | n factors | rank | MAC overhead | reconstruction error | |---|---|---|---|---| | 32,48 | 32,48 | 177 | 9.2x | 0.740 | 16,96 | 12,128 | 47 | 3.3x | 0.672 | Asymmetric shapes win on both axes at once — 2.8x less arithmetic and better accuracy. Every top-scoring configuration used n = 12,128 ; the output factorization dominates. This choice is free and worth ~3x. 1536 x 1536 layer Qwen2.5-1.5B q proj at 4x compression. Medians of 6 runs with a 20 s idle between them; ranges in brackets. Cold cache — a 32 MB buffer is read between iterations to evict weights from the 16 MB L2, simulating a real 28-layer forward pass where nothing stays resident. This is the regime that matters. | method | cold µs | vs TT-AMX | memory | |---|---|---|---| Dense FP32 cblas sgemv | 142.8 128–156 | 2.78x slower | 9.44 MB | | Dense FP16 BNNS | 102.0 68–111 | 1.99x slower | 4.72 MB | TT-AMX FP32 ours | 51.3 47–58 | 1.00x | 2.36 MB | INT4 group-64 estimated | ~26 | 0.5x — faster | 1.18 MB | Hot cache — reported because a reviewer will ask, and because it is the one regime where we lose: | method | hot µs | vs TT-AMX | |---|---|---| | Dense FP32 | 52.2 | 1.96x slower | TT-AMX FP32 ours | 26.6 | 1.00x | | Dense FP16 BNNS | 13.6 | 0.51x — faster | Peak kernel throughput: 889–947 GFLOP/s, 62–66% of this chip's AMX ceiling measured AMX peak 1424 GFLOP/s; NEON peak, for contrast, is only 415 GFLOP/s across all 6 P-cores — a hand-written NEON kernel provably cannot reach parity . No PyTorch, no model download, no vendored dependencies. Just Accelerate. git clone https://github.com/yourusername/tt-amx.git cd tt-amx cmake -B build -DCMAKE BUILD TYPE=Release cmake --build build ctest --test-dir build --output-on-failure 3/3 ./build/bench cold cache hot vs cold, TT vs dense fp32/fp16 ./build/bench amx chain 81-factorization sweep, GFLOP/s ./build/bench neon peak NEON roofline anchor To pack real weights instead of the synthetic build fixture: python3 tools/tt packer.py --model Qwen/Qwen2.5-1.5B-Instruct --out build/q proj.qtensor y = x · W where W M×N is stored as two TT cores, M = m1·m2 , N = n1·n2 : step 1 C1 n1·r, m2 = A1 n1·r, m1 · B1 m1, m2 A1 = G1.transpose 1,2,0 ← the ONLY permute, done offline by the packer B1 = x, natural layout, untouched step 2 out n1, n2 = A2 n1, r·m2 · B2 r·m2, n2 A2 = C1 reinterpreted: n1·r, m2 row-major IS n1, r·m2 ← 0 bytes moved B2 = G2 reinterpreted: r, m2, n2 IS r·m2, n2 ← 0 bytes moved The rule that makes the chain close: the core is always the left operand its layout is ours to choose offline and the activation is always the right . The 216 KB intermediate lives in L2 across both GEMMs — it never reaches DRAM. It does not fit in the 128 KB L1D; claims to the contrary in early drafts of this work were wrong. Verified end-to-end against numpy.einsum on the un-permuted cores: tests/test amx e2e.cpp . INT4 is still better for frozen-weight inference , on both speed ~2x and accuracy 0.132 vs 0.672 reconstruction error at the same 4x compression. TT-AMX targets the case where you need a differentiable parameter space. No end-to-end quality claim. At 4x compression TT reconstruction error on real Qwen weights is 0.67–0.84, versus 0.83 for a random Gaussian matrix of the same shape. TT extracts little structure these weights actually have; a recovery fine-tune is mandatory before any perplexity number means anything. Kernel benchmarks are unaffected — throughput does not depend on core values. We lose on hot cache to dense FP16 0.51x . The win requires the working set to exceed L2. FP16 via BNNS does not work for this shape. Measured 64 GFLOP/s — 14x slower than the FP32 cblas chain. BNNS reaches AMX for a fat dense GEMV but falls off a cliff on the skinny, deep GEMMs a TT chain produces step 2 is m=12, n=128, k=4512 . This is a measured negative, not future work. Single layer, batch 1. Batching makes TT worse : dense GEMV is memory-bound at batch 1 and consumes its spare compute for free as batch grows, while TT's MAC count scales linearly. Measured 1.65x → 26.6x penalty from batch 1 to 128. Absolute latencies are thermally sensitive. Sustained benchmarking on this machine shifted cold-cache TT latency from 43 µs to 74 µs. Ratios are far more stable than absolutes; quote ratios. tools/tt packer.py safetensors → TT-SVD → transpose 1,2,0 → .qtensor + golden vector include/engine/ format contract, mmap loader, AMX scheduler API src/engine/ mmap loader, AMX scheduler two sgemm calls bench/ implemented microbenchmarks roofline, sweep, cold cache tests/ C++ correctness suite + the Python analyses behind every number archive/ earlier scopes: full-LLM runtime, NEON kernel, superseded skeletons FINDINGS.md the full measurement protocol, sections A–J THIRD PARTY.md attribution; nothing is vendored or linked FINDINGS.md is the lab notebook — every claim above traces to a runnable script. | claim | script | |---|---| | MAC overhead, optimal contraction order | tests/test contraction cost.py | | machine roofline, NEON vs AMX ceilings | tests/test roofline.py , bench/bench neon peak.c | | batch ≥ 8 makes TT worse; size crossover | tests/test three paths.py | | TT accuracy vs SVD vs INT4 on real weights | tests/test tt validation.py | | factorization sweep, GFLOP/s | bench/bench amx chain.c | | hot vs cold cache, fp16 rejection | bench/bench cold cache.c | Two measurement bugs found and documented during this work, both of which initially produced wrong conclusions: unsigned underflow in benchmark data generation i%13 -6 with size t i , and a memset -based cache thrash that compiled to non-temporal stores and evicted nothing. See FINDINGS.md §I.1. MIT — see LICENSE .