--- title: 'Open-TQ-Metal: Efficient LLM Long-Context Inference' url: https://www.emergentmind.com/topics/open-tq-metal type: topic --- # Open-TQ-Metal: Efficient LLM Long-Context Inference Open-TQ-Metal encompasses a suite of methodologies and system-level innovations aimed at enabling resource-efficient long-context inference for large language models (LLMs) on Apple Silicon, specifically through the fusion of compressed-domain (int4) attention directly in the computational graph. It marks the first practical demonstration of 128,000-token context inference for dense 70B-parameter LLMs, such as Llama 3.1 70B, on a single 64GB consumer Mac—configurations previously unattainable due to KV cache memory bottlenecks and bandwidth constraints. The approach exploits on-the-fly KV cache quantization to 4-bit fixed-point and fused attention kernel design operating on the quantized representation, eliminating overhead from intermediate dequantization and dramatically compressing cache memory and inference latency while strictly maintaining accuracy equivalence to full-precision baselines [2604.16957]. ## 1. Motivation and Problem Setting Contemporary generative models exhibit inference working set footprints dominated by the key-value (KV) cache for persistent sequence memory, which scales linearly with sequence length $S$, number of layers $L$, head dimension $d$, and quantization width $b$: \[ \text{KV}_\text{bytes} = 2 \times L \times H_\text{kv} \times S \times d \times b \] For Llama 3.1 70B at 128K tokens, FP16 ($b=2$) yields a 40 GB KV cache, which in conjunction with other model parameters exceeds the 64 GB unified memory limit of consumer Apple Silicon (total: 79 GB). Existing frameworks—including mlx-lm, llama.cpp, and Ollama—are limited by the need to materialize FP16 KV matrices and perform explicit dequantization before every attention operation, incurring prohibitive memory and bandwidth demands, and failing at context lengths over 64–73K tokens. ## 2. Fused Compressed-Domain Attention: Algorithmic Framework Open-TQ-Metal introduces a custom single-pass "fused sdpa\_int4" kernel with the following workflow: - **On-the-fly Quantization:** Incoming KV vectors are quantized to 4 bits (int4), with group-wise (size $g=32$) scaling and zero-point correction. Two int4 values are packed per byte, achieving ≈3.2× compression. - **Fused Attention Kernel:** - Register-level dequantization is performed per key/value nibble using bitwise operations and group scale/zero metadata, with no staging of entire $S\times d$ matrices. - The attention kernel iterates over the compressed cache in a single pass. During each query (of shape $\mathbf q \in \mathbb R^d$), for every position $i$: 1. $\hat{\mathbf k}_i \gets \mathrm{dequant4}(\mathbf K_i)$, $\hat{\mathbf v}_i \gets \mathrm{dequant4}(\mathbf V_i)$ 2. $a_i \gets \mathbf q^\top \hat{\mathbf k}_i$ 3. Online softmax is computed and accumulated in-place: \[ m \leftarrow \max_i a_i,\quad \ell \leftarrow \sum_i e^{a_i - m},\quad \mathbf o \leftarrow \sum_i e^{a_i - m}\,\hat{\mathbf v}_i \] 4. Output is $\mathbf o/\ell$. - **No intermediate dequantization** or materialization: All operations are performed in the compressed (int4) domain, sidestepping substantial memory traffic. ## 3. Quantization Methodologies and Cross-Model Analysis The system implements group-wise int4 quantization: \[ s = \frac{\max(x_{g}) - \min(x_{g})}{15},\quad z = -\min(x_{g})/s,\quad q = \mathrm{clamp}\bigl(\mathrm{round}(x/s + z),0,15\bigr),\quad \hat x=(q-z)s \] With $g=32$, two values are packed per byte, and per-group $s, z$ metadata amortized accordingly. Comparison is made to angular (PolarQuant) quantization, where only the angle $\theta$ of each key vector is stored. The resulting perturbation on an attention score is proportional to the attention scale $\alpha$: \[ \Delta\mathrm{score} = \alpha \|\mathbf q\| \|\mathbf k\| (1-\cos\delta_\theta) \] For Gemma 4 (with $\alpha=1.0$) and Llama (with $\alpha\approx 1/\sqrt d$), identical $\delta_\theta$ results in up to $11\times$ more error in the former, leading to substantial accuracy degradation over deep stacks. The empirical finding is that the attention scale factor $\alpha$, not model size, is predictive of the viability of angular quantization: small $\alpha$ damps directional error (Llama), but large $\alpha$ (Gemma 4) amplifies it, rendering cosine-based quantizers unreliable outside low-$\alpha$ settings. ## 4. Metal Kernel Implementation and Apple Silicon Adaptation - **Compute Architecture:** Custom Metal compute shaders leverage Apple M1/M2 GPUs, using 32-wide SIMD groups to process per-head slices. MLX primitives and chained dispatches implement split-K parallelism for scaling KV cache length. - **Memory Management:** Keys and values are stored as packed uint8 arrays (2 int4 nibbles/byte); per-group metadata is amortized. This compresses the 40 GB FP16 KV cache at 128K tokens to 12.5 GB in int4. - **Split-K Parallelism:** The KV length is chunked into $C = \lceil S/512\rceil$ segments. Partial softmax operations are run in parallel, with an online softmax reduction to maintain $O(1)$ latency per decode step, even at maximal context. ## 5. Experimental Validation and Performance Benchmarks The system is profiled extensively on Llama 3.1 70B (80 layers, $d=128$, 8 KV heads) and Gemma 4 31B (60 layers). Key findings include: - **Kernel Speedup:** - At 128K context, the fused sdpa\_int4 achieves $9.9$ ms per attention, vs $480.6$ ms for the dequantize-then-attend baseline, a $48\times$ speedup. - At 64K and 16K contexts, observed speedups are $39\times$ and $20\times$, respectively. - **Memory Utilization:** | Context | FP16 KV Cache | int4 KV Cache | Total Memory (GB) | Fits in 64GB? | |----------|---------------|---------------|-------------------|---------------| | 1 K | 0.3 GB | 0.1 GB | 41.2 GB | ✓ | | 16 K | 5.0 GB | 1.6 GB | 42.7 GB | ✓ | | 64 K | 20.0 GB | 6.3 GB | 47.4 GB | ✓ | | 128 K | 40.0 GB | 12.5 GB | 53.6 GB | ✓ | - **End-to-End Throughput:** - Open-TQ-Metal: $6.0$ tok/s at 128K context (entire inference fits via int4 KV). - mlx-lm (FP16): $7.3$ tok/s but unable to exceed $\approx73$K tokens due to memory exhaustion. - Top-1 token predictions are bit-for-bit identical between fused int4 and full-precision baselines. ## 6. Quantizer Robustness Across Architectures Experiments on Gemma 4 31B and Llama 3.1 70B elucidate quantization error propagation: | Method | Llama $\alpha=0.0884$ | Gemma $\alpha=1.0$ | |-----------------------|-----------------------|---------------------| | PolarQuant 4b cosine | $\approx$ 0.958 | 0.621 | | PolarQuant 4b KL div | $\approx$ $8\!\times\!10^{-9}$ | $2\!\times\!10^{-4}$ | | Int4 (g=32) cosine | $\approx$ 0.998 | 0.992 | | Int4 (g=32) KL div | $<10^{-10}$ | $<10^{-8}$ | PolarQuant is effective for models with low $\alpha$ (e.g., Llama) but fails for standard or high $\alpha$ (e.g., Gemma 4), as angular error corrupts softmax attention across the stack. Int4 per-group quantization yields uncorrelated, nearly-cancelling errors effective across both architectures. ## 7. Impact and Future Directions Open-TQ-Metal establishes the first fused compressed-domain attention inference system on consumer Mac hardware, permitting full-length, large-parameter LLM inference without memory overflows and eliminating the tradeoff between context length and output fidelity. The framework challenges previous assumptions regarding quantizer–architecture interactions, highlighting the decisive role of the attention scale $\alpha$ in quantization error amplification and supporting an architectural co-design paradigm for future quantizer-hardware-model systems. Its rigorous evaluation spanning 330 experiments across Llama and Gemma families demonstrates the feasibility of 3.2× KV cache compression, up to 48× speedup, and strict preservation of decode accuracy—all in a deployment-ready implementation on Apple Metal [2604.16957]. Source: https://www.emergentmind.com/topics/open-tq-metal