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
Aug 20, 2026

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

Long-context modeling is a pivotal capability for Large Language Models, yet the quadratic complexity of attention remains a critical bottleneck, particularly during the compute-intensive prefilling phase. Our previous work, FlashPrefill, mitigates this cost through instantaneous pattern discovery and max-based dynamic thresholding; however, it remains an algorithmic prototype that is still distant from production deployment. In this paper, we present FlashPrefill V2, which evolves FlashPrefill from a prototype toward practical long-context serving along three dimensions. First, we introduce a mean correction term that effectively suppresses the approximation error, keeping performance degradation manageable even at extreme sparsity levels. Second, we redesign the sparse attention operator with PackGQA memory access, warp specialization, and pingpong pipelining, fully aligning with the latest FlashAttention-3/4 implementations and supporting FP8 inference to meet practical quantization requirements. Third, FlashPrefill V2 natively supports paged KV cache and continuous batching, allowing integration as an attention backend in modern inference frameworks such as SGLang. Extensive evaluations on NVIDIA H20 GPUs---among the most widely deployed inference accelerators---demonstrate that FlashPrefill V2 delivers up to 47.26x and 27.19x speedups over FlashAttention-2 at 128K context length under FP8 and BF16 precision, respectively, and, in FP8, still achieves a 30.49x speedup against an FA3/4-aligned dense baseline.

Qihang Fan, Huaibo Huang, Zhiying Wu +2

MAIS&NLPR, CASIA · UCAS · WeChat, Tencent

LLM Inference AccelerationLong-Context Language Model InferenceBlock-Sparse Attention
AskPDF
Oct 5, 2026

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

Fine-grained mixed-precision quantization promises efficient large language model inference, but local precision choices can conflict with regular GPU storage and computation units. This precision-boundary mismatch limits the translation of compression into practical acceleration. We introduce AlignQuant, a post-training quantization method that uses GPU-compatible two-dimensional weight tiles as the common unit of precision allocation, compact storage, and execution. This shared partition lets precision follow sensitivity within output channels. Joint prefill/decode calibration scores precision reductions using projection-output perturbations weighted by language-model loss gradients under quantized activations. Phase-normalized scores prioritize higher precision for tiles important to either phase under a model-wide weight-storage budget. Each tile stores one selected representation, while phase-specialized kernels reuse the packed model and expand lower-bit weights for INT8 computation with 8-bit activations. Across four LLMs spanning 3B to 14B parameters, AlignQuant achieves up to 2.50×2.50\times generation speedup over BF16 while preserving model quality. Evaluations further cover three GPUs and contexts up to 64K tokens. These results show that local precision flexibility and regular GPU execution can coexist through a shared tile unit. The implementation is available at https://github.com/HanzhiZhang-Ulrica/AlignQuant.

Hanzhi Zhang, Qiao Zhang, Qinglei Cao +4

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

Mixed-Precision QuantizationLLM QuantizationLLM Inference AccelerationPost-Training QuantizationNeural Network Quantization
AskPDF
Mar 3, 2026

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

The quadratic complexity of self-attention during the prefill phase impedes long-context inference in large language models. Existing sparse attention methods face a trade-off among context adaptivity, sampling overhead, and fine-tuning costs. We propose VSPrefill, a mechanism requiring lightweight training that uses the vertical-slash structural pattern in attention distributions. Our compact VSIndexer module predicts context-aware importance scores for vertical columns and slash diagonals from key-value representations augmented with RoPE. This approach constructs sparse masks with linear complexity without modifying the backbone parameters. During inference, an adaptive cumulative-threshold strategy allocates sparsity budgets per layer, while a fused kernel executes attention with on-the-fly index merging. Evaluated on Qwen3-4B-Instruct and LLaMA-3.1-8B-Instruct across the LongBench and RULER benchmarks, VSPrefill preserves 98.35% of the full attention accuracy while delivering a 4.95x average speedup at a context length of 128k. These results establish a new Pareto frontier in the trade-off between accuracy and efficiency.

Chen Guanzhong

School of Computer Science and Technology, Tongji University

Long-Context Language Model InferenceLLM Inference AccelerationSparse Attention
AskPDF