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

Transformer-based large language models increasingly rely on long contexts for tasks like document summarization, multi-page question answering, and retrieval-augmented generation. However, during the prefill phase, dense self-attention computes pairwise query-key scores leading to quadratic computational complexity and substantial memory traffic. Existing acceleration techniques typically apply a uniform low-precision format, select sparse token interactions, or use IO-aware fused attention kernels with fixed precision. These approaches leave spatial precision routing over attention score tiles outside the streaming loop. TileMix addresses this gap by treating numerical precision as an executable spatial decision over hardware-aligned score-tile groups within a fused dense attention kernel.

TileMix partitions the attention matrix into hardware-aligned compute tiles and organizes adjacent key tiles into routing groups governed by a grouping factor. For each query-tile row and KV head, the routing decisions are packed into a compact 64-bit bitmask word. Inside the latency-critical inner loop of the kernel, a constant-time shift-and-mask operation retrieves the precision decision for each score tile, dispatching it to either an FP16 Tensor Core path or an INT8 Tensor Core path with INT32 accumulation. After applying blockwise scale restoration, score tiles from both execution paths update a shared online-softmax state. This design preserves complete dense token connectivity, requires no model retraining, and natively supports grouped-query attention, variable-length batching, and INT8 key/value caches.

Evaluated across LongEval, LV-Eval, and A100 prefill benchmarks using model families such as LLaMA, Qwen, and Vicuna, TileMix demonstrates a controllable accuracy-efficiency frontier. By selectively routing a subset of score tiles to INT8 while keeping sensitive regions in FP16, TileMix recovers long-context quality that is otherwise lost under uniform INT8 attention while simultaneously improving prefill throughput compared to standard FP16 execution.