在苹果芯片上实现长文本推理,128K上下文仅需一台64GB Mac
Open-TQ-Metal: Fused Compressed-Domain Attention for Long-Context LLM Inference on Apple Silicon

- 直接在压缩数据上计算注意力,省去解压步骤
- 速度提升48倍,显存占用从40GB降至12.5GB
- 适合需要长文本处理的开发者和研究者
我们提出 Open-TQ-Metal,首个在 Apple Silicon 上实现融合压缩域注意力的框架,使单台64GB消费级 Mac 上运行 Llama 3.1 70B 的 128K 上下文推理成为可能。该方法实时将键值缓存量化为 int4,通过自定义 Metal 计算着色器直接在压缩表示上执行注意力计算,完全消除中间解压矩阵。在涵盖两个模型族(Gemma 4 31B 与 Llama 3.1 70B)的330次实验中,融合 sdpa_int4 内核在128K上下文下相比解压后计算基准实现48倍速度提升,键值缓存内存从40GB降低至12.5GB(压缩3.2倍),且保持与FP16推理相同的首位词预测结果。我们还首次开展跨架构的键值缓存量化分析,发现注意力缩放因子决定极化量化等方案成败:当 Gema 4 的 attn_scale=1.0 时,方向误差被放大25-100倍,远超使用标准 1/sqrt(d) 缩放的 Llama。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。