TileMix：以Tile為中心的混合精度注意力機制加速LLM推論

摘要

大型語言模型（LLM）中的長上下文預填充會產生可觀的計算與記憶體流量，因為密集自注意力會計算二次方的查詢-鍵分數。現有方法要麼採用均勻低精度路徑，要麼選擇性地處理詞元互動，均未在融合密集注意力內對硬體對齊的分數分塊進行空間精度路由。我們提出 TileMix，一種以分塊為中心的精度路由核心，將數值精度轉化為融合密集注意力中對分數分塊群組的可執行空間決策。TileMix 將注意力矩陣劃分為硬體對齊的分數分塊，將路由決策打包成緊湊的位元遮罩，並透過 FP16 或 INT8 分數計算分派每個分塊群組，同時兩條路徑皆更新共享的線上 softmax 狀態。可擴展的精度分組使每個路由位元能管轄多個相鄰的鍵分塊，從而在長上下文下維持硬體對齊的計算分塊與緊湊的中繼資料。透過對所有合法分塊群組進行路由，TileMix 保持密集的詞元連通性，無需訓練，並支援分組查詢注意力、可變長度批次與 INT8 鍵/值快取。在 LongEval、LV-Eval 與 A100 預填充基準上，針對 LLaMA、Qwen 與 Vicuna，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.