🤖 AI

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

TileMix is a training-free, tile-centric mixed-precision kernel that accelerates long-context LLM inference by dynamically routing hardware-aligned attention score tiles between FP16 and INT8 paths within a fused dense attention operation, thereby recovering the accuracy lost by uniform low-precision methods while improving throughput.

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

Published 2026-08-19

📖 5 min read🧠 Deep dive

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Large language models have become the engines behind a new generation of artificial intelligence, capable of summarizing entire books, answering complex questions from long documents, and holding conversations that span thousands of words. To do this, these models rely on a mechanism called attention, which allows them to weigh the importance of every word in a sentence against every other word. When the text is short, this process is quick. But as the context grows to include entire chapters or legal contracts, the computational cost explodes. The model must calculate a score for every possible pair of words, creating a massive grid of data that demands enormous amounts of memory and processing power. This bottleneck has made it difficult to run these powerful tools efficiently on standard hardware, especially when dealing with the long sequences required for real-world tasks like legal analysis or medical record review.

Researchers have tried to solve this by simplifying the math. One common approach is to lower the precision of the numbers the computer uses, switching from high-accuracy calculations to faster, lower-precision ones. This is like switching from measuring ingredients with a laboratory scale to using a kitchen spoon; it is much faster, but if you do it for every single step of a complex recipe, the final dish might taste wrong. Another approach is to skip calculations entirely, ignoring words that seem unimportant. While this saves time, it risks cutting off the very connections the model needs to understand the story. The challenge has been finding a way to use the speed of low-precision math without sacrificing the accuracy of the high-precision kind, all while keeping the model's ability to see the full picture.

A team of researchers has introduced a new method called TileMix that tackles this problem by treating the attention process not as a single, uniform block of work, but as a collection of smaller, manageable tiles. Imagine the massive grid of word-pair scores as a large mosaic. Instead of painting the entire mosaic with one type of paint, TileMix divides it into small, hardware-aligned square sections. For each of these sections, the system makes a split-second decision: does this specific group of word pairs need the high-precision, slow calculation, or can it get away with the fast, low-precision one? This decision is made based on a pre-planned map that groups these tiles together, allowing the computer to switch between high and low precision within the same operation without stopping to reorganize the data.

The core of this innovation is a routing system that acts like a traffic controller for the computer's processor. It packs these decisions into a compact code, essentially a string of bits that tells the processor which path to take for each tile. If a tile is marked for high precision, the processor uses its most accurate math units. If it is marked for low precision, it switches to its fastest, energy-efficient units. Crucially, both paths update a shared memory state that keeps track of the overall results, ensuring that the final output remains consistent. This allows the system to preserve the full connectivity of the model—meaning no word interactions are ever ignored—while still reaping the speed benefits of low-precision math for the parts of the calculation that can tolerate it.

In their experiments, the researchers tested this method on several popular large language models, including LLaMA, Qwen, and Vicuna, using sequences ranging from 16,000 to 64,000 words. They compared their approach against the standard high-precision method and a version that used low-precision math for everything. The results showed that using low-precision math for the entire calculation often led to a significant drop in accuracy, causing the model to miss details or fail at retrieval tasks. However, TileMix was able to recover most of that lost quality. By carefully routing only specific groups of tiles to the high-precision path, the system maintained accuracy levels very close to the full high-precision baseline, while simultaneously achieving much faster processing speeds.

The study also explored different patterns for deciding which tiles should get the high-precision treatment. They found that the arrangement mattered; some patterns, which kept high-precision calculations in specific spatial regions of the word grid, worked better than others depending on the task. For instance, certain layouts that preserved high precision for local word interactions performed better on factual recall tasks, while others were more robust for general question answering. The researchers demonstrated that this method could be applied without retraining the models, meaning it could be deployed immediately on existing systems. They also showed that the approach works well with variable-length batches and different model architectures, suggesting it is a flexible tool for improving efficiency.

Ultimately, TileMix offers a controllable balance between speed and accuracy. It suggests that the future of efficient long-context processing does not require choosing between being fast or being smart. Instead, by intelligently mixing the two, systems can process long documents at a speed that approaches the theoretical limits of the hardware, without the severe accuracy penalties that come from using low-precision math everywhere. This approach provides a practical path forward for deploying large language models in scenarios where both speed and precision are critical, such as real-time analysis of massive datasets or interactive tools that need to understand long, complex contexts instantly.

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

Problem Statement

Long-context inference in Large Language Models (LLMs) faces a significant bottleneck during the prefill phase due to the quadratic computational complexity ( $O(L^2)$ ) of dense self-attention. Existing acceleration strategies generally fall into three categories, each with limitations:

Low-Precision Quantization: Methods like INT8 quantization improve arithmetic and memory efficiency but typically apply a uniform precision path across the entire attention matrix. This often leaves spatial precision routing outside the fused attention loop, and uniform low-precision execution can degrade long-context quality.
Sparsity-Based Methods: Approaches that select active token interactions (e.g., sliding windows, sparse patterns) reduce computation but alter the dense connectivity of the attention graph, potentially losing critical token interactions.
IO-Aware Fused Attention: Kernels like FlashAttention partition attention into hardware-aligned tiles to optimize data movement and fuse operations (score computation, online softmax, value aggregation). However, these kernels traditionally maintain a uniform score-computation path, missing the opportunity to route specific spatial regions to different precision levels within a single execution.

The core challenge addressed is how to integrate heterogeneous precision (FP16 and INT8) into a single fused attention kernel without breaking the dense token connectivity or incurring significant metadata overhead, while ensuring that different arithmetic paths converge correctly to a shared online-softmax state.

Methodology: TileMix

TileMix introduces a tile-centric precision-routing kernel that treats numerical precision as a spatial decision over groups of score tiles within a fused dense attention computation.

Core Mechanisms

Tile-Group Precision Routing:
- The attention matrix ( $L_q \times L_k$ ) is partitioned into hardware-aligned compute tiles of size $BLOCK_M \times BLOCK_N$ .
- Instead of assigning precision per tile, TileMix groups adjacent key tiles along the key dimension. A single routing bit governs a group of $g$ key tiles (where $g$ is a grouping factor).
- This allows the kernel to dispatch each tile group to either an FP16 path (standard matrix multiplication) or an INT8 path (using INT8 Tensor Core MMA with INT32 accumulation) within the same streaming loop.
Compact Bitmask Routing:
- Routing decisions are packed into 64-bit integers (bitmasks), where each bit corresponds to a key-tile group.
- This enables constant-time lookup inside the inner loop using shift-and-mask operations ( $R_{hk,m,g_j} = ((b_{hk,m} \gg g_j) \& 1)$ ).
- The metadata overhead is $O(H_k T_m)$ , where $H_k$ is the number of KV heads and $T_m$ is the number of query tiles, remaining compact even as context length grows.
Shared-State Heterogeneous Execution:
- Both the FP16 and INT8 paths update a shared online-softmax state (row-wise maximum, normalizer, and unnormalized output accumulator).
- The INT8 path computes $Q_8 K_8^\top$ , applies block scales, and rescales the result to enter a common floating-point domain before updating the shared state.
- This design preserves the dense attention graph (all legal token interactions are computed) while varying the arithmetic precision based on spatial routing policies.
Policy and Layouts:
- Routing policies are defined by static, data-free structured templates (e.g., Band, Global, Row-Random, Aligned Sparse, BigBird, Sparse Transformer).
- These templates determine which spatial regions of the attention matrix are routed to FP16 (high precision) and which to INT8 (low precision) based on a target INT8 coverage ratio (e.g., 25%, 50%, 75%).
- The system supports Grouped-Query Attention (GQA), variable-length batching, and INT8 key/value caches.

Key Contributions

Tile-Centric Precision Routing: The paper introduces tile-group precision as a spatial execution abstraction for fused dense attention, enabling fine-grained FP16/INT8 allocation across all legal token interactions without sparsifying the graph.
Shared-State Heterogeneous Score Execution: A fused kernel is designed that aligns FP16 and INT8 score paths to a common score domain, integrating them through a single online-softmax recurrence.
Compact and Scalable Kernel-Native Routing: The authors develop packed bitmask routing with constant-time inner-loop lookup and scalable precision grouping that maintains hardware-aligned compute tiles while supporting long contexts.
Practical Long-Context Inference: The implementation supports GQA, variable-length batching, and INT8 KV caches, validated across multiple model families (LLaMA, Qwen, Vicuna) and benchmarks.

Experimental Results

The evaluation was conducted on NVIDIA A100 GPUs using LLaMA 3.2 3B, Qwen 2/2.5 7B, and Vicuna 7B.

Accuracy Recovery: Uniform INT8 attention ("One") often degrades retrieval and question-answering performance significantly as context length increases. TileMix recovers this lost quality by retaining FP16 precision for selected tile groups. For instance, on the LongEval benchmark, mixed-routing configurations (e.g., SpTrans or BigBird layouts) consistently outperformed uniform INT8, often approaching or matching FP16 baselines while maintaining high INT8 coverage.
Throughput Efficiency: TileMix improves prefill throughput over standard FP16 FlashAttention. At 4k tokens, configurations like SpTrans75 (75% INT8) achieved ~31.8 K tokens/s compared to ~14.3 K tokens/s for FlashAttention. TileMix configurations generally outperformed or matched uniform INT8 baselines in end-to-end throughput, demonstrating that the overhead of routing and mixed-precision handling is negligible compared to the gains from INT8 Tensor Cores.
Numerical Stability: Analysis of mean absolute deviation from a fixed FP16 reference showed that deviation increases with INT8 coverage but remains controllable. The choice of routing layout (spatial distribution of FP16) significantly impacts accuracy, with certain layouts (like SpTrans) showing robustness across different models and tasks.
Heavy-Hitter Protection: Static routing layouts like SpTrans were shown to naturally protect high-importance attention interactions (heavy hitters) from being routed to INT8, routing only a small fraction of high-mass interactions to low precision despite high nominal INT8 coverage.

Significance and Claims

The paper claims that TileMix establishes score-tile-group precision routing as a kernel-native abstraction for long-context attention. Its primary significance lies in:

Decoupling Connectivity from Precision: It allows the preservation of dense token connectivity (unlike sparse attention) while enabling heterogeneous precision execution (unlike uniform quantization).
Controllable Accuracy-Efficiency Frontier: It provides a practical mechanism to trade off accuracy and efficiency by adjusting the INT8 coverage ratio and routing layout, offering a "controllable accuracy-efficiency frontier" across model families.
Training-Free Deployment: The method operates without retraining, relying on static spatial templates to guide precision routing.
Kernel-Level Optimization: By integrating precision routing directly into the fused attention kernel with constant-time lookup, it avoids the latency penalties associated with invocation-level precision selection or external scheduling.

The authors conclude that spatial precision control is a practical and effective dimension for optimizing dense attention kernels, recovering long-context quality lost under uniform low-precision methods while significantly accelerating inference.