Search papers, labs, and topics across Lattice.
This paper introduces TileMix, a novel tile-centric precision-routing kernel designed to enhance the efficiency of long-context prefill in large language models (LLMs) by optimizing numerical precision across hardware-aligned score tiles. By partitioning the attention matrix and utilizing compact bitmasks for routing decisions, TileMix enables the use of mixed precision (FP16 and INT8) without sacrificing token connectivity or requiring additional training. The method demonstrates significant improvements in prefill throughput and quality across various benchmarks, establishing a controllable accuracy-efficiency frontier for different model families.
TileMix achieves a breakthrough in LLM inference by enabling mixed-precision attention that boosts throughput while maintaining long-context quality.
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.
Tsinghua AI