通过分块精度路由,让大模型长文本推理更快更准。
TileMix: Tile-Centric Mixed-Precision Attention for LLM Inference Acceleration

- 将注意力矩阵分块,按需选择FP16或INT8计算每组块。
- 在不损失长文本质量的前提下,比FP16快3.2倍,比统一INT8准15%。
- 无需训练,支持多种模型和动态长度输入,适合部署优化。
长上下文预填充阶段因密集自注意力计算二次方查询-键得分而带来巨大计算与内存开销。现有方法或采用统一低精度路径,或选择性跳过交互,未能对硬件对齐的得分块进行精细化精度调度。本文提出TileMix,一种以块为中心的精度路由内核,在融合的密集注意力中将数值精度作为可执行的空间决策。该方法将注意力矩阵划分为硬件对齐的得分块,将路由决策压缩为紧凑位掩码,并使每组块通过FP16或INT8计算,同时共享在线Softmax状态。可扩展的精度分组设计使得每个路由位控制多个相邻键块,保持硬件对齐计算块并减少元数据开销。通过路由所有合法块组,TileMix保留了密集令牌连接,无需训练,支持分组查询注意力、变长批量及INT8键值缓存。在LLaMA、Qwen、Vicuna的LongEval、LV-Eval和A100预填充基准上,TileMix恢复了统一INT8下的长上下文质量损失,相比FP16提升预填充吞吐量,实现了跨模型家族的可控精度-效率权衡。代码已开源。
原文摘要 · Abstract (English)
Long-context prefill in large language models (LLMs) incurs substantial computation and memory traffic because dense self-attention computes quadratic query-key scores. Existing methods either use a uniform low-precision path or select token interactions, leaving spatial precision routing over hardware-aligned score tiles outside fused dense attention. We introduce TileMix, a tile-centric precision-routing kernel that makes numerical precision an executable spatial decision over score-tile groups within fused dense attention. TileMix partitions the attention matrix into hardware-aligned score tiles, packs routing decisions into compact bitmasks, and dispatches each tile group through FP16 or INT8 score computation while both paths update a shared online-softmax state. Scalable precision grouping lets each routing bit govern multiple adjacent key tiles, preserving hardware-aligned compute tiles and compact metadata at long contexts. By routing all legal tile groups, TileMix preserves dense token connectivity, requires no training, and supports grouped-query attention, variable-length batches, and INT8 key/value caches. Across LongEval, LV-Eval, and A100 prefill benchmarks on LLaMA, Qwen, and Vicuna, TileMix recovers long-context quality lost under uniform INT8 and improves prefill throughput over FP16, yielding a controllable accuracy-efficiency frontier across model families. The implementation is available at https://github.com/HanzhiZhang-Ulrica/TileMix.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。