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

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:

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.

Hot cache — reported because a reviewer will ask, and because it is the one regime where we lose:

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 adifferentiableparameter 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 arandom Gaussian matrixof 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 FP32cblas

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 TTworse: 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 , AMX scheduler API
src/engine/               mmap , 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.

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

.