{"slug": "tt-amx-a-zero-copy-tensor-train-inference-engine-for-apple-silicon", "title": "TT-AMX, a zero-copy Tensor-Train inference engine for Apple Silicon", "summary": "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.", "body_md": "A bare-metal C++ engine that runs **Tensor-Train (TT/MPO) compressed linear layers**\non Apple Silicon by dispatching the contraction onto the undocumented **AMX\ncoprocessor**. At 4x compression it beats a dense FP16 GEMV by **~2x under\nrealistic cold-cache conditions**, using half the memory.\n\nTensor-Train is attractive for on-device work because it is a **continuously\ndifferentiable** compression manifold — unlike discrete INT4 quantization — which\nmakes it usable for on-device PEFT, continual learning, and quantum-inspired ML.\nHistorically it has been unusable for a different reason: a 5–10x latency\npenalty.\n\n**That penalty turns out to be arithmetic, not cache.** A TT matvec does 3–35x\nmore multiply-accumulates than the dense matvec it replaces (measured with\nnumpy's *optimal* contraction order). No amount of cache tuning changes a MAC\ncount. This engine wins anyway, by making those extra MACs nearly free.\n\n**1. Zero-copy ahead-of-time core permutation.** A single offline\n`transpose(1,2,0)`\n\non core 1 turns the runtime into two back-to-back\n`cblas_sgemm`\n\ncalls with *nothing between them* — no loop, no permute, no copy.\nThe intermediate is reinterpreted, not moved.\n\n**2. Asymmetric factorization.** Balanced tensor shapes are the wrong default.\nSweeping 81 factorization pairs on real Qwen2.5-1.5B weights:\n\n| m factors | n factors | rank | MAC overhead | reconstruction error |\n|---|---|---|---|---|\n| (32,48) | (32,48) | 177 | 9.2x | 0.740 |\n(16,96) |\n(12,128) |\n47 |\n3.3x |\n0.672 |\n\nAsymmetric shapes win on **both** axes at once — 2.8x less arithmetic *and*\nbetter accuracy. Every top-scoring configuration used n = (12,128); the output\nfactorization dominates. This choice is free and worth ~3x.\n\n`1536 x 1536`\n\nlayer (Qwen2.5-1.5B `q_proj`\n\n) at 4x compression. Medians of 6 runs\nwith a 20 s idle between them; ranges in brackets.\n\n**Cold cache** — a 32 MB buffer is read between iterations to evict weights from\nthe 16 MB L2, simulating a real 28-layer forward pass where nothing stays\nresident. *This is the regime that matters.*\n\n| method | cold µs | vs TT-AMX | memory |\n|---|---|---|---|\nDense FP32 (`cblas_sgemv` ) |\n142.8 [128–156] | 2.78x slower | 9.44 MB |\n| Dense FP16 (BNNS) | 102.0 [68–111] | 1.99x slower | 4.72 MB |\nTT-AMX FP32 (ours) |\n51.3 [47–58] |\n1.00x |\n2.36 MB |\nINT4 group-64 (estimated) |\n~26 | 0.5x — faster |\n1.18 MB |\n\n**Hot cache** — reported because a reviewer will ask, and because it is the one\nregime where we lose:\n\n| method | hot µs | vs TT-AMX |\n|---|---|---|\n| Dense FP32 | 52.2 | 1.96x slower |\nTT-AMX FP32 (ours) |\n26.6 |\n1.00x |\n| Dense FP16 (BNNS) | 13.6 | 0.51x — faster |\n\nPeak kernel throughput: **889–947 GFLOP/s, 62–66% of this chip's AMX ceiling**\n(measured AMX peak 1424 GFLOP/s; NEON peak, for contrast, is only 415 GFLOP/s\nacross all 6 P-cores — a hand-written NEON kernel provably cannot reach parity).\n\nNo PyTorch, no model download, no vendored dependencies. Just Accelerate.\n\n```\ngit clone https://github.com/yourusername/tt-amx.git\ncd tt-amx\ncmake -B build -DCMAKE_BUILD_TYPE=Release\ncmake --build build\nctest --test-dir build --output-on-failure    # 3/3\n./build/bench_cold_cache    # hot vs cold, TT vs dense fp32/fp16\n./build/bench_amx_chain     # 81-factorization sweep, GFLOP/s\n./build/bench_neon_peak     # NEON roofline anchor\n```\n\nTo pack real weights instead of the synthetic build fixture:\n\n```\npython3 tools/tt_packer.py --model Qwen/Qwen2.5-1.5B-Instruct --out build/q_proj.qtensor\n```\n\n`y = x · W`\n\nwhere `W (M×N)`\n\nis stored as two TT cores, `M = m1·m2`\n\n, `N = n1·n2`\n\n:\n\n```\nstep 1   C1(n1·r, m2) = A1(n1·r, m1) · B1(m1, m2)\n         A1 = G1.transpose(1,2,0)   ← the ONLY permute, done offline by the packer\n         B1 = x, natural layout, untouched\n\nstep 2   out(n1, n2)  = A2(n1, r·m2) · B2(r·m2, n2)\n         A2 = C1 reinterpreted: (n1·r, m2) row-major IS (n1, r·m2)   ← 0 bytes moved\n         B2 = G2 reinterpreted: (r, m2, n2)         IS (r·m2, n2)    ← 0 bytes moved\n```\n\nThe rule that makes the chain close: **the core is always the left operand**\n(its layout is ours to choose offline) and **the activation is always the\nright**. The 216 KB intermediate lives in L2 across both GEMMs — it never\nreaches DRAM. (It does not fit in the 128 KB L1D; claims to the contrary in\nearly drafts of this work were wrong.)\n\nVerified end-to-end against `numpy.einsum`\n\non the un-permuted cores:\n`tests/test_amx_e2e.cpp`\n\n.\n\n**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`\n\nchain. 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.\n\n```\ntools/tt_packer.py        safetensors → TT-SVD → transpose(1,2,0) → .qtensor + golden vector\ninclude/engine/           format contract, mmap loader, AMX scheduler API\nsrc/engine/               mmap loader, AMX scheduler (two sgemm calls)\nbench/                    implemented microbenchmarks (roofline, sweep, cold cache)\ntests/                    C++ correctness suite + the Python analyses behind every number\narchive/                  earlier scopes: full-LLM runtime, NEON kernel, superseded skeletons\nFINDINGS.md               the full measurement protocol, sections A–J\nTHIRD_PARTY.md            attribution; nothing is vendored or linked\n```\n\n`FINDINGS.md`\n\nis the lab notebook — every claim above traces to a runnable script.\n\n| claim | script |\n|---|---|\n| MAC overhead, optimal contraction order | `tests/test_contraction_cost.py` |\n| machine roofline, NEON vs AMX ceilings | `tests/test_roofline.py` , `bench/bench_neon_peak.c` |\n| batch ≥ 8 makes TT worse; size crossover | `tests/test_three_paths.py` |\n| TT accuracy vs SVD vs INT4 on real weights | `tests/test_tt_validation.py` |\n| factorization sweep, GFLOP/s | `bench/bench_amx_chain.c` |\n| hot vs cold cache, fp16 rejection | `bench/bench_cold_cache.c` |\n\nTwo measurement bugs found and documented during this work, both of which\ninitially produced *wrong* conclusions: unsigned underflow in benchmark data\ngeneration (`(i%13)-6`\n\nwith `size_t i`\n\n), and a `memset`\n\n-based cache thrash that\ncompiled to non-temporal stores and evicted nothing. See `FINDINGS.md`\n\n§I.1.\n\nMIT — see `LICENSE`\n\n.", "url": "https://wpnews.pro/news/tt-amx-a-zero-copy-tensor-train-inference-engine-for-apple-silicon", "canonical_source": "https://github.com/ansarzeinulla/tensor-train-amx", "published_at": "2026-08-22 14:55:01+00:00", "updated_at": "2026-08-22 15:14:14.550758+00:00", "lang": "en", "topics": ["artificial-intelligence", "machine-learning", "ai-infrastructure", "ai-research"], "entities": ["TT-AMX", "Apple Silicon", "AMX coprocessor", "Qwen2.5-1.5B", "GitHub"], "alternates": {"html": "https://wpnews.pro/news/tt-amx-a-zero-copy-tensor-train-inference-engine-for-apple-silicon", "markdown": "https://wpnews.pro/news/tt-amx-a-zero-copy-tensor-train-inference-engine-for-apple-silicon.md", "text": "https://wpnews.pro/news/tt-amx-a-zero-copy-tensor-train-inference-engine-for-apple-silicon.txt", "jsonld": "https://wpnews.pro/news/tt-amx-a-zero-copy-tensor-train-inference-engine-for-apple-silicon.jsonld"}}