cs.AIAug 18, 2026

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

Authors: Hanzhi Zhang, Qiao Zhang, Qinglei Cao, Heng Fan, Yan Huang, Kewei Sha, Yunhe Feng

Organizations: LLaVi Lab, Computer Science & Engineering, University of North Texas · Computer Science, Saint Louis University · Data Science, University of North Texas

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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.

Figures & tables

Figure 1: Comparison of attention efficiency strategies. Colors denote execution states: blue for high precision, gray for low precision, and white for removed interactions. Quantization reduces weight/activation memory but often keeps attention softmax and accumulation in higher precision. Sparse/block attention executes a selected subset of token interactions through structured or dynamic patterns. TileMix preserves all legal token interactions and routes score-tile groups through FP16 or INT8 paths inside one fused attention kernel.
Figure 2: TileMix fused attention with tile-group precision routing. Left: The Lq×LkL_{q}\times L_{k} attention matrix is processed in BLOCKM×BLOCKN\mathrm{BLOCK}_{M}\times\mathrm{BLOCK}_{N} compute tiles. For each query-tile row mm (outer loop), the kernel streams key/value tiles nn (inner loop) from HBM to on-chip SRAM and registers, dispatches the score tile Sm n=QmKn⊤/dS_{m}^{\,n}=Q_{m}K_{n}^{\top}/\sqrt{d} to either FP16 matmul or INT8 Tensor Core MMA with INT32 accumulation and rescaling, and updates a shared FP16 online-softmax state to produce OmO_{m} without materializing the full score matrix. Right: A binary routing map is grouped along the key-tile dimension with width BLOCKNmask=g BLOCKN\mathrm{BLOCK}_{N}^{\text{mask}}=g\,\mathrm{BLOCK}_{N} and packed per (hk,m)(h_{k},m) into a 64-bit mask bhk,mb_{h_{k},m} . The bit at position gj=⌊n/g⌋g_{j}=\lfloor n/g\rfloor selects the arithmetic path for the corresponding key-tile group (1: INT8, 0: FP16). Kernel-side legality masks, including causality and boundary conditions, are enforced independently of the routing map.
Figure 3: Key-tile grouping and bitmask encoding for one query-tile row mm . Adjacent compute tiles along the key dimension form routing groups gjg_{j} , each controlled by one bit in the packed mask bhk,mb_{h_{k},m} (1: INT8, 0: FP16).
Figure 4: Constant-time shift-and-mask routing lookup. Shifting bit gjg_{j} to the least-significant position yields the arithmetic-path decision Rhk,m,gjR_{h_{k},m,g_{j}} for the corresponding key-tile group.
Figure 5: Line-level retrieval accuracy on the LongEval benchmark for LLaMA 3.2 3B under different tile-group routing layouts, evaluated across prompt lengths from 3.1k to 38.7k tokens. Line plots (left y-axis) report exact-match retrieval accuracy, while bar plots (right y-axis, Δ\Delta accuracy vs. FP16) show differences relative to the FP16 attention baseline. Each panel corresponds to a routing layout with 25%, 50%, or 75% of tile groups routed to INT8. Colors indicate routing settings: green denotes the original FP16 attention baseline, red denotes One (all legal score-tile groups routed to INT8) , and blue denotes mixed configurations with partially INT8-routed tile groups.
Table 1: LV-Eval long-context question answering accuracy for LLaMA 3.2 3B . We compare FP16 attention, One (100% INT8) , representative TileMix routing layouts, sparse long-context baselines (MInference and FlexPrefill), and an INT8 attention baseline (SageAttention). Sparse baselines reduce computation by pruning token interactions, whereas TileMix preserves dense legal connectivity and routes legal tile groups to different score-computation paths. Cells are color-coded by comparison to One and FP16: red indicates performance below One , blue indicates performance between One and FP16, and green indicates performance above FP16.
Seq Len |0%5%10%
1k |7.27×10−57.27{\times}10^{-5}7.47×10−47.47{\times}10^{-4}1.19×10−31.19{\times}10^{-3}
2k |5.89×10−55.89{\times}10^{-5}4.29×10−44.29{\times}10^{-4}1.16×10−31.16{\times}10^{-3}
4k |4.87×10−54.87{\times}10^{-5}6.71×10−46.71{\times}10^{-4}1.10×10−31.10{\times}10^{-3}
8k |6.84×10−56.84{\times}10^{-5}8.34×10−48.34{\times}10^{-4}6.32×10−36.32{\times}10^{-3}
Seq Len |15%20%25%
1k |1.56×10−31.56{\times}10^{-3}1.67×10−31.67{\times}10^{-3}2.03×10−32.03{\times}10^{-3}
Table 3: Mean absolute output deviation from the fixed Torch FP16 reference under different score-tile-group INT8 coverage ratios; lower values indicate closer agreement with the reference implementation.
Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

CategorySymbolDescription
AttentionQ,K,VQ,K,VQuery, key, and value matrices: Q∈RLq×dQ\in\mathbb{R}^{L_{q}\times d} , K∈RLk×dK\in\mathbb{R}^{L_{k}\times d} , V∈RLk×dV\in\mathbb{R}^{L_{k}\times d}
Lq,LkL_{q},L_{k}Query and key lengths; Lq=Lk=LL_{q}=L_{k}=L for self-attention
ddQuery, key, and value head dimension
SSScaled score matrix: S=QK⊤/d∈RLq×LkS=QK^{\top}/\sqrt{d}\in\mathbb{R}^{L_{q}\times L_{k}}
PPAttention weights: P=softmax(S)∈RLq×LkP=\mathrm{softmax}(S)\in\mathbb{R}^{L_{q}\times L_{k}}
OOAttention output: O=PV∈RLq×dO=PV\in\mathbb{R}^{L_{q}\times d}
Appendix
Table 4: Notation for attention, tiling, online softmax, quantization, head mapping, and routing.
Figure 6: Visualization of attention layouts across layers and heads. The figure compares the attention values from a LLaMA 3.2 3B Instruct model on the Multi-News dataset.
Figure 8: Line-level retrieval accuracy on the LongEval benchmark for Qwen 2.5 7B under different tile-group routing layouts, evaluated across prompt lengths from 3.1k to 38.7k tokens. Line plots (left y-axis) report exact-match retrieval accuracy, while bar plots (right y-axis, Δ\Delta accuracy vs. FP16) show differences relative to the FP16 attention baseline. Each panel corresponds to a routing layout with 25%, 50%, or 75% of tile groups routed to INT8. Colors indicate routing settings: green denotes the original FP16 attention baseline, red denotes One (all legal score-tile groups routed to INT8), and blue denotes mixed configurations with partially INT8-routed tile groups.
Figure 9: Line-level retrieval accuracy on the LongEval benchmark for Qwen 2 7B under different precision policy layouts, evaluated across prompt lengths from 3.1k to 38.7k tokens.
Table 12
Table 13
Table 14
Seq Len |Torch Value |FlashAttn |100% INT8 |75% INT8 |50% INT8 |25% INT8
1024 |3.02 ⁣× ⁣1003.02\!\times\!10^{0} |3.55 ⁣× ⁣10−13.55\!\times\!10^{-1} |3.56 ⁣× ⁣10−13.56\!\times\!10^{-1} |2.22 ⁣× ⁣10−12.22\!\times\!10^{-1} |1.14 ⁣× ⁣10−11.14\!\times\!10^{-1} |5.86 ⁣× ⁣10−35.86\!\times\!10^{-3}
2048 |2.90 ⁣× ⁣1002.90\!\times\!10^{0} |3.18 ⁣× ⁣10−13.18\!\times\!10^{-1} |4.48 ⁣× ⁣10−14.48\!\times\!10^{-1} |2.96 ⁣× ⁣10−12.96\!\times\!10^{-1} |1.37 ⁣× ⁣10−11.37\!\times\!10^{-1} |5.00 ⁣× ⁣10−35.00\!\times\!10^{-3}
4096 |3.13 ⁣× ⁣1003.13\!\times\!10^{0} |3.79 ⁣× ⁣10−13.79\!\times\!10^{-1} |3.69 ⁣× ⁣10−13.69\!\times\!10^{-1} |2.51 ⁣× ⁣10−12.51\!\times\!10^{-1} |1.28 ⁣× ⁣10−11.28\!\times\!10^{-1} |5.86 ⁣× ⁣10−35.86\!\times\!10^{-3}
8192 |3.21 ⁣× ⁣1003.21\!\times\!10^{0} |3.31 ⁣× ⁣10−13.31\!\times\!10^{-1} |3.98 ⁣× ⁣10−13.98\!\times\!10^{-1} |2.60 ⁣× ⁣10−12.60\!\times\!10^{-1} |1.34 ⁣× ⁣10−11.34\!\times\!10^{-1} |5.86 ⁣× ⁣10−35.86\!\times\!10^{-3}
Appendix
Table 17: Numerical behavior of a 32-layer attention model on random inputs. Torch Value reports the maximum absolute logit magnitude under the selected Torch FP16 reference; all other columns report maximum absolute deviation from this reference.
Table 16
Table 20: Direct comparison between FP16 and FP32 accumulation under different precision layouts and mixing ratios.

Explore similar work

CardsList

  1. FlashPrefill V2: Block-Sparse Prefill Attention for Long-Context LLM Serving

    Aug 20, 2026Qihang Fan, Huaibo Huang, Zhiying Wu +2LLM Inference AccelerationLong-Context Language Model Inference

    MAIS&NLPR, CASIA · UCAS · WeChat, Tencent

  2. AlignQuant: Tile-Aligned Mixed-Precision Quantization for Efficient LLM Generation

    Oct 5, 2026Hanzhi Zhang, Qiao Zhang, Qinglei Cao +4Mixed-Precision QuantizationLLM Quantization

    LLaVi Lab, Computer Science & Engineering, University of North Texas

  3. VSPrefill: Vertical-Slash Sparse Attention with Lightweight Indexing for Long-Context Prefilling

    Mar 3, 2026Chen GuanzhongLong-Context Language Model InferenceLLM Inference Acceleration

    School of Computer Science and Technology, Tongji University