TileMix: Tile-Centric Mixed-Precision Attention for LLM Inference Acceleration

2026-08-17

AI Summary

为解决长上下文预填充中的计算和内存流量问题，TileMix通过在融合密集注意力中使用基于瓦片的精度路由内核来加速LLM推理。

Abstract

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.

Problem

• long-context prefill
• large language models
• dense self-attention
• numerical precision
• hardware-aligned score tiles

Innovation

• Tile-Centric Precision-Routing
• Fused Dense Attention
• Long-Context Prefill
• Scalable Precision Grouping
• Hardware-Aligned Score Tiles