Introducing Open-TQ-Metal: Open-Source Fused Attention for Long-Context LLM Inference

Apr 21, 2026 | Sai Vegasena

Introducing Open-TQ-Metal - an open-source, Metal-native implementation of fused compressed-domain attention, extending Google Research's TurboQuant approach (Zandieh et al., ICLR 2026) to Apple Silicon. It enables Llama 3.1 70B at 128K context on a single 64GB Mac, a configuration none of the major open-source inference frameworks (mlx-lm, llama.cpp, Ollama) can reach on this hardware.

Llama 3.1 70B at 128K context needs 79 GB of memory. Normally, a 64 GB Mac can't do it. Open-TQ-Metal writes fused Metal attention kernels that compress the KV cache from 40 GB to 12.5 GB. Now 128K fits - and runs at 6 tok/s.

This is the first implementation of fused compressed-domain attention on Apple Silicon, built across 330 experiments on two model families (Gemma 4 31B and Llama 3.1 70B).

Paper: arxiv.org/abs/2604.16957

Key Contributions

• First Metal implementation of fused compressed-domain attention. Custom Metal compute shaders compute attention directly on int4-quantized KV cache data, with zero intermediate dequantization buffers. 48x speedup at 128K context over the dequantize-then-attend baseline.
• Extended the TurboQuant approach from 8B to 70B. The original TurboQuant paper tested on 8B models with 32 layers. We ran 330 experiments across Gemma 4 31B (60 layers) and Llama 3.1 70B (80 layers), identifying which techniques transfer and which break at depth.
• A cross-architecture quantization finding. A single architectural parameter (attn_scale) determines whether angular quantization methods like PolarQuant succeed or fail - a finding that was not in the original TurboQuant, PolarQuant, or QJL papers.
• Fully open-sourced. Metal shader code, C++ inference engine, and all benchmarks are public.

Overview

Standard int4 inference dequantizes the full key matrix to FP32 before computing attention - wasting bandwidth and memory on temporary matrices that grow with context length:

The fused sdpa_int4 Metal kernel reads packed int4 keys and values directly from device memory, dequantizes per-element via bitwise operations in GPU registers, and computes attention with online softmax - no intermediate matrices ever materialize.

Split-K parallelism divides long KV sequences into 512-token chunks processed in parallel across Metal's GPU cores, then reduces partial softmax results via chained MLX Primitives. This solves the Metal dispatch race condition (multiple kernels within a single eval_gpu() execute without barriers) and achieves flat latency scaling from 1K to 128K tokens.

Key Discoveries

The attn_scale Finding

The most surprising result: the attention scale factor - not model size - determines whether angular quantization methods work.

PolarQuant encodes KV vectors as recursive polar coordinates. It works beautifully on Llama (attn_scale = 1/√d = 0.0884) but produces gibberish on Gemma 4 (attn_scale = 1.0). The reason: Gemma's undampened scale amplifies a 4% directional error 25-100x more per layer. Across 60 layers, this compounds through softmax into incoherent output.

Int4 per-group quantization works on both architectures because its per-element affine errors are uncorrelated and cancel in dot products - unlike PolarQuant's structured angular bias.

Wired Memory Is Mandatory

A 39 GB model on 64 GB hardware causes macOS to page weights to SSD. One line fixes a 12x slowdown: mx::set_wired_limit(max_working_set_size) → 0.6 → 6.0 tok/s.

Results

Kernel Speedup (Llama 3.1 70B)

Constant Throughput (Gemma 4 31B)

The fused kernel maintains constant 10 tok/s regardless of context length. The baseline degrades from 10.8 to 7.2 tok/s as dequantization bandwidth increases.

Memory: 128K Context on 64 GB

End-to-End

Open-TQ-Metal trades 18% throughput for 3.2x context capacity - the only framework that runs Llama 70B at 128K on consumer hardware.

Memory breakdown at each context length. The dashed red line is the 64 GB limit. At 128K, FP16 KV needs 79.1 GB (overflows). Open-TQ-Metal fits in 51.6 GB.

What Didn't Work (and Why)

Paper Comparison

The original TurboQuant paper (ICLR 2026) proposes PolarQuant + QJL on CUDA for 8B models. Open-TQ-Metal ports it to Metal, scales to 70B, and discovers that int4 with fused attention outperforms the paper's pipeline on real hardware.

Run It Yourself

The inference engines are in separate repos, each with full build instructions, pretrained weight loaders, and streaming chat interfaces:

Both are C++ and Metal with Python only for tokenization. Clone, build MLX from source, cmake && make, and chat.

How We Discovered These Breakthroughs: The Ensue Agent Swarm

The 330-experiment sweep across Gemma 4 31B and Llama 3.1 70B was done with our agent swarm and coordinated over a shared memory layer, Ensue.

A swarm of orchestrator agents worked in parallel conducting sub agents on separate Macs. Each relied on custom retrieval systems sifting through memories under one data hierarchy. The memories were indexed through insights, hypotheses, task delegation, and reasoning. The aim was to asynchronously conduct an efficient sweep of the optimization frontier with diverse, creative experiments.

Our sub agents were categorized under hardware profiler, data diagnostician, model builder, experiment runner, plateau analyst, validator, and reflector. Each provided separate reasoning traces with complementary agentic loops that enhanced the capability of the overall swarm. The cross-experiment synthesis that connects findings across architectures, parameter sweeps, and differing philosophies culminates to an enriched collective intelligence.

Optimizing ML Systems

The kind of work that produced Open-TQ-Metal kernel-level optimization, cross-architecture analysis, and agent-coordinated experiment sweeps is the kind of ML systems work we do for clients.

Check out our previous work on how we compressed months of ML work to 48 hours and fixed LLM inference slowdown and other case studies.

If you're hitting memory walls, scaling a paper's approach to your stack, or trying to ship inference optimizations, book a call or contact us.