TileMix：面向大语言模型推理加速的以磁贴为中心的混合精度注意力机制

摘要

大型语言模型（LLMs）中的长上下文预填充（prefill）会带来大量计算与内存流量，因为密集自注意力需要计算数量呈二次方规模的查询-键分数。现有方法要么使用统一的低精度路径，要么选择性地交互token，而未能在融合密集注意力内部对硬件对齐的分数块进行空间精度路由。我们提出TileMix，一种以块为中心的精度路由内核，它将数值精度转化为融合密集注意力中分数块组上的可执行空间决策。TileMix将注意力矩阵划分为硬件对齐的分数块，将路由决策打包为紧凑的位掩码，并通过FP16或INT8分数计算调度每个块组，而两条路径同时更新共享的在线softmax状态。可扩展的精度分组使每个路由位能够管理多个相邻的键块，从而在长上下文中保持硬件对齐的计算块和紧凑的元数据。通过对所有合法块组进行路由，TileMix保留了密集的token连接性，无需训练，并支持分组查询注意力、变长批次和INT8键/值缓存。在LLaMA、Qwen和Vicuna上的LongEval、LV-Eval和A100预填充基准测试中，TileMix恢复了在统一INT8下丢失的长上下文质量，并相较于FP16提高了预填充吞吐量，在多个模型家族中提供了可控的精度-效率前沿。实现代码见 https://github.com/HanzhiZhang-Ulrica/TileMix。

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.