Open-TQ-Metal: Fused Compressed-Domain Attention for Long-Context LLM Inference on Apple Silicon

Abstract: We present Open-TQ-Metal, the first implementation of fused compressed-domain attention on Apple Silicon, enabling 128K-context inference for Llama 3.1 70B on a single 64GB consumer Mac -- a configuration impossible with all existing inference frameworks. Open-TQ-Metal quantizes the KV cache to int4 on the fly and computes attention directly on the compressed representation via custom Metal compute shaders, eliminating all intermediate dequantization matrices. Across 330 experiments spanning two model families (Gemma 4 31B and Llama 3.1 70B), the fused sdpa_int4 kernel achieves 48x attention speedup at 128K context over the dequantize-then-attend baseline, reduces KV cache memory from 40 GB to 12.5 GB (3.2x compression), and maintains identical top-1 token predictions to FP16 inference. We further provide the first cross-architecture analysis of KV cache quantization methods, revealing that the attention scale factor -- not model size -- determines whether angular quantization schemes like PolarQuant succeed or fail, with Gemma 4's attn_scale=1.0 amplifying directional error 25-100x more than Llama's standard 1/sqrt(d) scaling.

• The paper introduces a fused int4 attention mechanism that eliminates intermediate dequantization matrices, drastically reducing memory consumption for long-context LLMs.
• It achieves a 48× speedup and a 3.2× reduction in memory overhead compared to conventional FP16 KV cache methods on Apple Silicon.
• The study emphasizes the role of attention scaling in mitigating quantization errors, offering insights for hardware-efficient transformer designs.

Open-TQ-Metal: Fused Compressed-Domain Attention for Long-Context LLM Inference on Apple Silicon

Introduction

Open-TQ-Metal presents a fusion of compressed-domain attention computation with Metal-optimized GPU kernels, targeting ultra-long-context LLM inference on Apple Silicon. The core technical advance is a fully fused int4 attention mechanism that obviates intermediate dequantization matrices, pushing context capacity on consumer-grade Mac hardware past the physical VRAM wall set by previous frameworks. The paper integrates systematic quantization method analysis across diverse model architectures and substantiates all claims with extensive empirical results, shedding light on the interaction between transformer scaling factors and quantization error amplification.

Motivation and Background

KV cache size grows linearly with context length and model layers, becoming the dominant memory consumer during inference for large transformer models at high sequence lengths. For Llama~3.1~70B, the FP16 KV cache alone requires 40 GB for 128K tokens, exceeding the 64 GB unified memory of top-tier Apple Silicon. Existing frameworks hit this limit due to non-fused quantization and the cost of intermediate full-precision matrices. Recent works on KV cache quantization (KIVI, KVQuant, GEAR, TurboQuant, PolarQuant, QJL) focus on compression but fall short in efficiently combining quantization with SIMD-parallel attention computation on Metal.

Methodology

Open-TQ-Metal's core technical innovation is a fused int4 Scaled Dot-Product Attention (SDPA) kernel implemented as a Metal compute shader. This kernel reads packed int4 KV tensors from memory, applies on-the-fly bitwise dequantization (using vectorized nibble extraction and affine scaling), and executes attention in registers with online softmax—completely eliminating large dequantized matrices.

A split-K approach divides long context sequences into chunks, processed in parallel and then reduced via chained MLX Primitives, maintaining latency scalability. A hybrid routing selects kernel execution modes for prefill (multi-query, MLX SDPA with quantize-store sequence) and decode (single query, fused kernel from int4 cache).

Figure 1: Llama~3.1~70B at 128K context—FP16 KV cache (79.1 GB) exceeds device capacity; int4 KV reduces total to 53.6 GB, enabling feasibility on 64 GB hardware.

Memory and Throughput Results

Empirically, the fused kernel achieves a 48$\times$ speedup over the dequantize-then-attend baseline at 128K context in Llama~70B, with sublinear latency scaling due to register-based dequantization and parallel chunking. The memory overhead is reduced by 3.2$\times$, enabling long-context (~236K tokens) inference on hardware previously limited to ~73K tokens with FP16.

Figure 2: Gemma~4~31B decode throughput vs. context length—fused kernel sustains constant throughput while the baseline suffers as dequantization bandwidth increases.

Figure 3: Memory breakdown for Llama~3.1~70B—Open-TQ-Metal (left) enables 128K context at 53.6 GB, mlx-lm (right) fails at 80 GB.

FP16 KV cache configurations in other frameworks (mlx-lm, llama.cpp, Ollama) cannot instantiate the 128K context window with 70B models on Apple Silicon. Open-TQ-Metal is the only stack that achieves this, albeit with a moderate (18%) throughput penalty due to dynamic cache management—this does not affect kernel performance itself.

Quantization Method Analysis

A rigorous cross-architecture analysis exposes how attention scale ($\alpha$) critically modulates quantization-induced degradation. PolarQuant, which quantizes vector angles, performs well on Llama ($\alpha=1/\sqrt{d}\approx0.088$, many layers) but fails on Gemma~4 ($\alpha=1.0$, different normalization scheme). With Gemma~4's scaling, angular errors are not suppressed and aggregate destructively through layers, leading to output collapse even at moderate context, whereas Llama’s scale factor dampens the error by 25–100$\times$.

Figure 4: KV cache size at 128K context for Llama~70B—int4 and PolarQuant 5b achieve practical sizes, but QJL fails due to noise accumulation in deep networks.

QJL achieves high compression but is only viable for shallow models—there is an exponential dropoff in signal with depth (correlation $\rho\approx0.85$ per layer, so $0.85^{80}\to0$ at 70B/80 layers). Hybrid strategies (partial QJL) offer minimal incremental gain (<2%).

Hardware and Engineering Considerations

The Metal kernel is engineered to match Apple Silicon's SIMD organization, with design choices tightly coupled to memory architecture and scheduling specifics of MLX/Metal. Wiring model weights (via mx::set_wired_limit) prevents paging-induced throughput collapse. Pre-allocated buffer management remains an engineering gap, as per-step KV concatenation introduces overhead not fundamental to the algorithm.

Figure 5: Total memory vs. context length for Llama~70B—FP16 KV crosses the 64 GB device limit at ~73K tokens, whereas int4 enables up to 236K sequences.

Figure 6: Gemma~4~31B at 256K context—only int4 compression enables model fitting within the 64 GB memory envelope of the M1 Max.

Figure 7: Standalone kernel latency on Gemma~4~31B—fused kernel (orange) drastically outperforms the MLX baseline (teal) as context grows.

Practical Limitations and Implications

The algorithm’s irreducible throughput ceiling is set not by the fused kernel, but by the bandwidth required to broadcast model weights (e.g., 98ms/token for Llama~70B, given 39 GB/400 GBps), highlighting the need for future exploration of MoE- or weight-compressed architectures. Int4 quantized KV cache can produce identical top-1 predictions at moderate contexts, but upper bounds on perplexity and extreme-context generalization (e.g., RULER, long-context QA) require further investigation, especially for architectures with non-standard attention scaling.

The sensitivity of angular quantization to scale factors sets an actionable precedent: quantization method selection must be co-designed with model scaling and normalization strategies. The engineering insights—particularly fused in-register dequant-atten, split-K on computation graphs, and sensitivity to attention scale—are directly portable to CUDA and other backends, but kernel reimplementation is non-trivial.

Conclusion

Open-TQ-Metal validates the practicality of fused compressed-domain attention for consumer hardware LLM inference, delivering large-context capacity and computation efficiency unattainable in existing Metal stacks. It exposes architectural determinants of quantization success, especially the decisive role of attention scaling in error propagation, and makes headway towards portable yet hardware-efficient LLM deployment. Future directions include extending these methods to more heterogeneous architectures, deeper empirical evaluation of extreme-context quality, and development of weight-compressed/active parameter architectures to break the bandwidth wall on consumer accelerators.