# TT-AMX, a zero-copy Tensor-Train inference engine for Apple Silicon > Source: > Published: 2026-08-22 14:55:01+00:00 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` .