TileMix: Tile-Centric Mixed-Precision Attention for LLM Inference Acceleration
Hanzhi Zhang, Qiao Zhang, Qinglei Cao, Heng Fan, Yan Huang, Kewei Sha, Yunhe Feng
cs.AI
2026-08-18
TileMix routes score tiles to FP16 or INT8 without dropping connections. LLaMA 3.2 3B on A100 hits 31.80K tokens/s at 4k prefill, about 2.2× FlashAttention, quality near full FP16.
Long-context prefill is dominated by dense self-attention. Score computation is O(L²) in sequence length, so the kernel burns both FLOPs and HBM traffic. Existing accelerators pick one of three levers. Quantized attention (SageAttention and related INT8 kernels) usually binds a single arithmetic path per launch. Sparse methods (MInference, FlexPrefill, Longformer-style patterns) drop token interactions. FlashAttention tiles the work for IO, but still runs one precision through the streaming loop.
The missing lever is spatial precision: keep every legal connection, and let each hardware-aligned score tile pick FP16 or INT8. That is TileMix.
TileMix is a training-free fused kernel from LLaVi Lab at the University of North Texas and Saint Louis University. Precision becomes an inner-loop dispatch over score-tile groups.
The Lq×Lk score matrix is cut into BLOCKM × BLOCKN compute tiles, the same SRAM units FlashAttention uses. Adjacent key tiles form a routing group. One bit per group selects the path: 1 for INT8, 0 for FP16. Those bits pack into a 64-bit word per KV head and query-tile row. A grouping factor g stretches one word across long keys so the number of groups stays at most 64. Lookup is a shift-and-mask. Metadata scales as O(Hk × Tm).
Think of painting a chessboard by squares. Fine pen on some squares, coarse pen on others. No piece is removed.
The INT8 path quantizes Q and K only, blockwise and symmetric: 128×d query blocks, 64×d key blocks, absmax over 127. Tensor Cores run INT8 MMA into INT32, then restore with block scales and 1/√d. V and the PV update stay FP16. After rescaling, both paths enter one floating-point score domain and update a shared FP16 online-softmax state: row max, normalizer, output accumulator. Causality and boundary legality are enforced separately from the routing map.
Evaluated policies are static, data-free templates borrowed from sparse-attention geometry: Band (diagonal neighborhood), Global, Row-Random, Aligned Sparse, BigBird as the union, and Sparse Transformer (SpTrans) with stride plus tail. Reported 25/50/75% figures are INT8 shares of legal tile groups, not of FLOPs. The same template is reused across layers, batch items, and KV heads. The kernel supports grouped-query attention, variable-length batches, and an INT8 KV-cache interface for decode. The implementation is Triton, autotuned once on A100 and then frozen.
The lead model is LLaMA 3.2 3B, with Vicuna-7B, Qwen-2-7B, and Qwen-2.5-7B in the appendix. Hardware is NVIDIA A100 40GB. Baselines: dense FP16, uniform INT8 on the same kernel (One), FlashAttention, MInference, FlexPrefill, and SageAttention. Throughput is end-to-end prefill at batch 8, including quantization, scale restore, routing, and scheduling.
LV-Eval long-context QA on LLaMA 3.2 3B, one representative slice:
| Method | 16k | 32k | 64k |
| FP16 | 32.04 | 15.08 | 7.75 |
| One (all INT8) | 28.78 | 11.62 | 5.42 |
| SpTrans25 | 31.75 | 15.56 | 8.01 |
| SageAttention | 29.79 | 13.50 | 5.77 |
| MInference | 26.93 | 11.55 | 5.31 |
| FlexPrefill | 27.11 | 11.74 | 5.39 |
Uniform INT8 drops. Mixed routing returns near FP16, and sometimes slightly above. Sparse baselines, which prune edges, sit lower on most subsets. On LongEval line retrieval, layout matters more than the nominal INT8 budget: rowrand and sptrans hold up as coverage rises; alignsparse, band, and global look better under conservative coverage.
The paper also reports SpTrans far above FP16 on 16k factual recall across LLaMA and Qwen. On LLaMA that subset, FP16 is 6.72 and SpTrans25 is 21.04. That jump does not look like ordinary quantization noise.
Prefill throughput on LLaMA 3.2 3B-Instruct, K tokens/s:
| Length | FlashAttention | One | SpTrans75 | SageAttention |
| 1k | 17.45 | 32.27 | 33.50 | 19.91 |
| 4k | 14.33 | 29.80 | 31.80 | 19.91 |
| 8k | OOM | 27.41 | 26.61 | 18.79 |
At 4k, SpTrans75 is 31.80 against FlashAttention's 14.33, about 2.22×, and close to all-INT8. At 8k under this protocol, PyTorch and FlashAttention OOM; TileMix still runs. Mean absolute deviation from Torch FP16 is about 10⁻⁵ at 0% INT8, 6.32×10⁻³ at 8k with 10% coverage, and 6.84×10⁻³ at 25%. Appendix SpTrans25 sends only about 8.5% of high-importance attention mass through INT8, below its 25% tile-group label.
Inference stacks get a new knob: precision can be a spatial assignment over hardware tiles, instead of a single choice among full FP16, uniform INT8, and dropped edges. No retraining. GQA and variable-length batches are already in the kernel. Deployment cost is low if the hardware has INT8 Tensor Cores.
This is incremental systems work, not a new attention algorithm. Layout and coverage have to be chosen per model and task. End-to-end timings include the whole pipeline, so a high-coverage mix can beat One when dispatch and memory layout cooperate.
The authors list three. Evaluation is prefill. The kernel is FP16/INT8 Tensor Cores on A100; FP8 or INT4 would need new scales and schedules. Routing in the tables is static; the interface can take adaptive maps, but those maps are not compared.
A few more caveats sit in the numbers. Quality headlines use 3B, with 7B in the appendix and no 70B-class run. Throughput stops at 8k while quality goes to 64k. Coverage is tile groups, not FLOPs. The 16k factual-recall spike over FP16 is large enough to look like a template-benchmark coupling, and it needs an independent check. Static routing does not find heavy hitters online, so a domain or model shift can hurt. Decode has an INT8 KV-cache hook; the published timings are almost all prefill.