PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2608.17336v1cs.AI

TL;DR

Long-context prefill is bottlenecked by quadratic dense-attention computation and memory traffic. TileMix routes FP16 or INT8 arithmetic across score-tile groups within fused attention, recovering quality lost under uniform INT8 while improving prefill throughput over FP16.

  • Problem

    Dense self-attention creates O(L^2) score computation and substantial intermediate-memory traffic, making long-context attention a primary inference bottleneck.

  • Method

    TileMix routes each legal score-tile group to FP16 or INT8 computation within fused dense attention, combining packed bitmasks with a shared online-softmax state.

  • Results

    Across LongEval, LV-Eval, prefill benchmarks, and numerical analyses, TileMix recovers long-context quality from uniform INT8 while improving prefill throughput over FP16.

  • Takeaways & Limitations

    TileMix establishes a controllable accuracy-efficiency frontier by adding spatial precision control to dense attention kernels.

  • Takeaways & Limitations

    The evaluated implementation targets long-context prefill on NVIDIA A100 GPUs with FP16 and INT8 Tensor Core paths; other formats require format-specific scale handling and kernel scheduling.

Abstract

from arXiv · show

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.

1 Introduction

TileMix addresses the O(L2) prefill bottleneck of dense self-attention by routing hardware-aligned score-tile groups between FP16 and INT8 within one fused execution. Its packed, scalable routing and shared online-softmax state preserve dense token connectivity while supporting practical long-context inference settings.

  • 1 Introduction: Dense self-attention computes interactions between all query and key tokens, creating O(L2) score computation and a primary execution bottleneck for long-context workloads.Long-context use cases include document summarization, multi-page question answering, and retrieval-augmented generation.
  • 1 Introduction: A fused FlashAttention-style kernel decodes packed routing words with constant-time bit operations, dispatches FP16 or INT8 Tensor Core score paths, and merges both through shared online softmax.After INT8 rescaling, both paths update the shared row-wise maximum, normalizer, and output accumulator.
  • 1 Introduction: TileMix introduces tile-group precision routing for dense attention, enabling fine-grained FP16/INT8 allocation across all legal token interactions.The approach preserves dense connectivity rather than removing token interactions, while making precision a spatial execution decision over score tiles.
  • 1 Introduction: Packed bitmask routing has O(HkTm) metadata, while scalable precision grouping lets one routing bit govern adjacent key tiles without abandoning hardware-aligned compute tiles.This design targets compact, regular execution at long contexts.
  • 1 Introduction: TileMix supports grouped-query attention, variable-length batching, and INT8 key/value caches for practical long-context inference evaluation.The paper validates these capabilities through long-context retrieval, question answering, and prefill-efficiency evaluation.

2 Related Work

Transformer acceleration primarily uses low-precision quantization, IO-aware tiled attention, and structural sparsity to reduce long-context inference cost through numerical compression, data-movement optimization, or selective token connectivity. Existing methods commonly fix precision at coarse tensor, operator, stage, or kernel levels, while sparse methods reduce computation by selecting token interactions.

  • 2 Related Work: Transformer acceleration follows three main directions: low-precision quantization, IO-aware tiled attention, and structural sparsity.These approaches target long-context inference cost through numerical compression, data-movement optimization, or selective token connectivity.
  • 2 Related Work: Quantization reduces memory and arithmetic costs with formats such as INT8 and INT4, while calibration, activation transformation, outlier handling, and blockwise scaling improve robustness.These designs typically assign formats at the tensor, operator, attention-stage, or quantization-block level.
  • 2 Related Work: IO-aware kernels process SRAM-resident tiles and fuse score computation, online softmax, and value aggregation, avoiding full attention-matrix materialization in HBM.The tiles govern data movement and work partitioning, while arithmetic precision is commonly fixed per kernel invocation or attention stage.
  • 2 Related Work: Sparse attention reduces computation by selecting token interactions through local-global patterns, content-based selection, clustering, positional mechanisms, adaptive architectures, or dynamically constructed active blocks.Examples include sliding-window, strided, dilated, and mixed local-global patterns.

3 Preliminaries

TileMix builds on FlashAttention’s tiled online-softmax execution and blockwise quantization to route FP16 and INT8 score tiles through different arithmetic paths while updating one shared streaming state. This addresses the O(L^2) intermediate-memory footprint and HBM traffic of materialized long-context attention without materializing the full score matrix.

  • 3 Preliminaries: Self-attention forms S = QK^⊤, P = softmax(S), and O = PV, while materializing S and P incurs O(L^2) intermediate memory and substantial HBM traffic.These costs grow with sequence length L and motivate tiled execution.
  • 3 Preliminaries: FlashAttention streams key/value tiles, keeps score and probability tiles on chip, and maintains row-wise maxima and normalizers for online softmax before writing normalized outputs.For each query tile, the kernel iterates over key tiles and divides the accumulated output by the final normalizer.
  • 3 Preliminaries: Blockwise quantization represents operand blocks with low-precision values and scales; TileMix instantiates INT8 to use A100 INT8 Tensor Cores with INT32 accumulation.Aligning quantization blocks with attention tiles enables efficient application of the corresponding scales.
  • 3 Preliminaries: Mixed-precision attention is a shared-state kernel problem because FP16 and INT8 score tiles differ numerically yet must enter a common domain before updating the same maximum, normalizer, and output accumulator.TileMix dispatches each score tile to FP16 matmul or INT8 MMA with rescaling, then updates a shared FP16 online-softmax state without materializing the full score matrix.
  • 3 Preliminaries: Routing groups adjacent key tiles into packed 64-bit masks, where each bit selects INT8 or FP16 computation while legality masks independently enforce causality and boundary conditions.A group spans BLOCKmask_N = g BLOCKN key-tile width, and bit position ⌊n/g⌋ selects the corresponding arithmetic path.

4 TileMix: Tile-Group Precision Routing

TileMix integrates FP16 and INT8 score computation into dense fused attention by routing hardware-aligned tile groups through packed bitmasks and constant-time inner-loop dispatch. Both arithmetic paths update a shared online-softmax state, while grouping preserves compute-tile alignment and supports long sequences.

  • 4 TileMix: Tile-Group Precision Routing: TileMix combines per-tile-group FP16/INT8 routing, packed bitmasks, and fused attention so heterogeneous score arithmetic executes without retraining.Each query-tile row reads routing decisions from packed words, and legal tile groups select the arithmetic path inside the streaming inner loop.
  • 4 TileMix: Tile-Group Precision Routing: Grouping lets one routing bit control adjacent key tiles, extending compact routing words across long sequences while preserving hardware-aligned BLOCKM × BLOCKN compute tiles.A group of g tiles spans g · BLOCKN key tokens, while the compute tile remains the execution unit.
  • 4 TileMix: Tile-Group Precision Routing: The routing interface supports grouped-query attention and variable-length batches, sharing decisions across query heads mapped to one KV head and using prefix-sum metadata for padding-free execution.For GQA, query heads mapped to the same KV head share tile-group routing decisions.
  • 4 TileMix: Tile-Group Precision Routing: Packed 64-bit masks provide constant-time shift-and-mask lookup with O(HkTm) routing metadata consumed directly by the streaming kernel.The routing map is indexed by KV head, query-tile row, and key-tile group.

5 Experiments

TileMix is evaluated for long-context retrieval and question answering, A100 prefill efficiency, and numerical behavior, primarily on LLaMA 3.2 3B with extensions to Vicuna and Qwen models. Results show improved throughput, recoverable long-context quality under routed precision, and controllable numerical deviation while preserving dense legal connectivity.

  • Experimental Setup: Experiments cover LV-Eval question answering, LongEval exact-match retrieval, and A100 prefill efficiency, using dense FP16, uniform INT8, sparse, and INT8 attention baselines.The main evaluation uses LLaMA 3.2 3B; appendices extend quality and efficiency studies to Vicuna-7B, Qwen-2-7B, and Qwen-2.5-7B.
  • Long-context Question Answering: TileMix narrows uniform INT8’s long-context QA quality gap by assigning FP16 to selected score-tile groups while preserving all legal interactions.Table 1 compares FP16, One, TileMix layouts, sparse baselines, and SageAttention; uniform INT8 often trails FP16 on LV-Eval.
  • Long-context Retrieval: Retrieval quality depends on FP16 placement as well as INT8 coverage: row_rand and sptrans remain stronger at higher coverage, while align_sparse, band, and global retain quality conservatively.Figure 5 evaluates LLaMA 3.2 3B from 3.1k to 38.7k tokens at 25%, 50%, and 75% INT8 tile-group coverage.
  • Prefill Efficiency: 31.80 K tokens/s: SpTrans75 at 4k tokens exceeds FlashAttention’s 14.33 K tokens/s and One’s 29.80 K tokens/s in end-to-end prefill.Measurements include quantization, scale restoration, routing, memory staging, and scheduling costs on identical A100 40GB hardware with batch size 8.
  • Numerical Behavior: Numerical deviation generally increases with INT8 coverage, with a marked rise between 5% and 10% at 8k tokens before remaining at a similar scale through 25%.The 0% configuration remains close to the fixed Torch FP16 reference, making coverage a practical numerical-control knob.

6 Conclusion

TileMix establishes score-tile-group precision routing as a kernel-native abstraction for long-context attention, varying FP16 and INT8 precision within dense attention while preserving the attention graph. Its training-free design recovers quality lost under uniform INT8, improves prefill throughput over FP16, and enables a controllable accuracy-efficiency frontier.

  • Conclusion: The kernel assigns FP16 and INT8 paths across legal tile groups, preserving dense attention connectivity while varying precision within one computation.Both paths update a shared online-softmax state, enabling mixed-precision execution without removing token interactions.
  • Conclusion: Scalable grouping spans adjacent key tiles while retaining hardware-aligned compute tiles and compact metadata through packed bitmasks and constant-time lookup.This makes score-tile-group precision routing executable inside the attention kernel.
  • Conclusion: TileMix recovers long-context quality lost under uniform INT8 while improving prefill throughput over FP16, establishing a controllable accuracy-efficiency frontier.The result demonstrates spatial precision control as a practical dimension for dense attention kernels.

Limitations

TileMix is evaluated for forward inference during long-context prefill, with the current implementation limited to FP16 and INT8 Tensor Core paths on NVIDIA A100 GPUs. Other numerical formats require format-specific scale handling and kernel scheduling, although the interface can accept static or adaptive routing policies.

  • TileMix targets forward inference during long-context prefill, the deployment setting evaluated throughout the work.
  • The current implementation instantiates tile-group routing with FP16 and INT8 Tensor Core paths on NVIDIA A100 GPUs.
  • Other numerical formats require format-specific scale handling and kernel scheduling.
  • Static routing templates provide deterministic policy construction, compact metadata, and constant-time kernel dispatch, while the kernel interface can consume alternative static or adaptive policies.

Ethical Considerations · A Notations

TileMix is presented as a systems method whose broader ethical impact is computational rather than data- or capability-related. Its notation is organized by functional roles matching attention computation, kernel design, quantization, and tile-group routing.

  • Ethical Considerations: The method’s primary broader impact is computational because it improves long-context LLM inference efficiency.
  • Ethical Considerations: TileMix does not introduce new training data, human-subject data, or model capabilities.
  • Ethical Considerations: Improved attention efficiency may reduce GPU time and energy consumption per supported workload.
  • Ethical Considerations: Lower inference cost could also increase aggregate deployment and total compute demand.
  • A Notations: Notation is grouped by functional role rather than order of appearance to support clarity and ease of reference.
  • A Notations: The notation organization follows the structure of attention computation and kernel design.
  • A Notations: Grouped notation covers core attention definitions, IO-aware tiling and indexing, online-softmax state variables, and blockwise quantization primitives.
  • A Notations: The notation also includes grouped-query attention head mappings and tile-group routing metadata used in the kernel execution flow.

B Precision Layouts · C Quantization Configuration · D Storage and Execution Layout

TileMix routes dense attention tile groups through FP16 or INT8 score paths without removing token interactions, using compact layouts, shared online softmax, and hardware-aligned execution. Its primary configuration quantizes Q/K while retaining FP16 V/PV computation, with quantize-once and on-the-fly operand preparation trading memory for data movement.

  • B Precision Layouts: TileMix converts spatial attention templates into hardware-aligned FP16/INT8 precision decisions while preserving dense token connectivity and separating precision routing from interaction legality.Causal and sequence-boundary masks independently determine legal interactions; routing layouts only select the arithmetic path for each legal tile group.
  • B Precision Layouts: 25%, 50%, and 75% settings denote legal score-tile-group INT8 coverage, with minor deviations caused by discrete routing groups and sequence boundaries.These percentages do not represent the fraction of total attention FLOPs executed in INT8.
  • B Precision Layouts: Layouts include uniform All INT8 and All FP16 references plus structured Band, Global, Row-Random, Aligned Sparse, BigBird, and SpTrans precision-routing templates.BigBird and SpTrans are the representative structured layouts reported in the main LV-Eval evaluation, while templates are statically constructed and reused rather than selected from benchmark outputs.
  • C Quantization Configuration: The primary score-routing configuration applies blockwise symmetric INT8 quantization to Q and K, retains FP16 V/PV computation, and merges both score paths into shared online softmax.It uses BLKQ = 128 and BLKK = 64, with FP32 intermediates for tile maxima, exponentiation, and reductions before updating FP16 state.
  • C Quantization Configuration: Quantize-once generates Q8, K8, and block scales once per attention call, reusing them across INT8 groups while preserving original FP16 Q/K for FP16 groups.An additional on-the-fly mode generates INT8 fragments only for routed groups and performs scale computation and quantization on chip; INT8 key/value caches are also supported separately.
  • D Storage and Execution Layout: Packed 64-bit routing words indexed by KV-head and query-tile row select each legal tile group’s score path while dense FlashAttention-style addressing and connectivity remain intact.Dense FP16 Q/K/V stay in contiguous HBM, while routing metadata is stored separately and read by the fused inner loop.
  • D Storage and Execution Layout: Fused on-the-fly matches FlashAttention’s 96 MB resident-HBM footprint, while TileMix uses 96 KB on-SM live state and a 32 KB accumulator versus 128 KB and 64 KB.Quantize-once instead uses 136.01 MiB resident HBM but reduces one INT8-routed QK tile’s operand read from 32 KB to 16.25 KB.

E Model Performance

Across Qwen, LLaMA, and Vicuna, uniform INT8 often loses long-context retrieval and QA accuracy, while TileMix mixed routing recovers quality by selectively retaining FP16 tile groups. The best layout and INT8 coverage remain model- and task-dependent, but dense token connectivity is preserved.

  • Long-context Retrieval: Mixed-routing configurations recover retrieval accuracy relative to uniform INT8 across Qwen 2.5 7B, Qwen 2 7B, and Vicuna 7B, especially at higher INT8 coverage.Uniform INT8 often degrades as prompt length increases, whereas mixed routing keeps selected tile groups in FP16 and routes the remainder to INT8.
  • Cross-model Trends: TileMix preserves dense token connectivity while recovering quality through tile-group precision routing, although the strongest layout and INT8 ratio depend on model and prompt length.The reported results identify bigbird, row_rand, and sptrans as the most consistent recovery layouts across evaluated models, particularly at higher INT8 coverage.
  • Long-context Question Answering: SpTrans mixed-precision configurations close much of the gap to FP16 on LV-Eval while outperforming or matching sparse baselines in most reported settings.This supports recovering long-context quality through precision routing rather than sparse token selection.
  • Long-context Question Answering: Across LV-Eval datasets and lengths, mixed routing usually narrows uniform INT8’s deficit, with moderate INT8 ratios often benefiting layouts whose tolerance for coverage varies.align_sparse, band, and global favor more conservative ratios, whereas bigbird, row_rand, and sptrans often tolerate higher INT8 coverage.

F Efficiency

TileMix provides implementation-level prefill efficiency gains across LLaMA, Qwen, and Vicuna, with mixed-routing configurations forming a controllable throughput–TOPS frontier. Higher INT8 coverage generally improves efficiency, and several mixed-routing configurations match or exceed the uniform INT8 reference while preserving executable score-tile routing.

  • F Efficiency: The efficiency evaluation compares Torch, FlashAttention, One’s uniform INT8 path, and TileMix mixed-routing variants across four models and multiple INT8 ratios.The reported ratios are 25%, 50%, and 75%, using batch size 8 with three warmup and five measurement iterations.
  • F Efficiency: Across models and sequence lengths, several mixed-routing configurations match or exceed One, while higher INT8 coverage generally increases throughput and TOPS.Tables 11–14 report throughput (Thpt, K tokens/s) and TOPS from 1k to 8k where executable; layout ordering depends on complete-pipeline execution.
  • F Efficiency: Torch has the lowest throughput and earlier OOM behavior, whereas FlashAttention is a strong FP16 baseline when memory permits.The comparison spans LLaMA 3.2 3B, Qwen 2.5 7B, Qwen 2 7B, and Vicuna 7B with model-dependent maximum sequence lengths.

G Numerical Analysis … L Max Mean | L Max Mean | L Max Mean | L Max Mean

TileMix’s numerical behavior is assessed against fixed Torch FP16 and FlashAttention references across depth, precision layouts, accumulation modes, larger models, and heavy-hitter routing. The results show that INT8 coverage and spatial layout dominate deviations, while FP16 accumulation remains stable and structured routing protects important attention mass.

  • G Numerical Analysis: The numerical analysis uses fixed Torch FP16 outputs as a common reference for deviations caused by quantization, accumulation order, rounding, and reduction scheduling.The appendix examines kernel deviation, depth and length, INT8 coverage, FlashAttention differences, accumulation precision, and heavy-hitter exposure.
  • G.1 Depth and Sequence-Length Effects: INT8 coverage and precision layout cause larger numerical deviations than accumulation precision or sequence length, with uniform INT8 deviating more than mixed routing.Across depths, deviations increase with model depth; reducing INT8-routed tile coverage generally reduces deviation from the fixed Torch FP16 reference.
  • G.2 Direct Comparison with FlashAttention: TileMix and FlashAttention differ even at 0% INT8 because their fused-kernel schedules differ, while MeanDiff generally increases with INT8 coverage.Repeated MaxDiff or MeanDiff values indicate equality at the reported precision rather than bitwise-identical outputs.
  • G.3 Pattern and Accumulation Effects: FP16 and FP32 accumulation produce similar deviations under matched precision layouts, so routing layout and INT8 coverage have the larger numerical effects.The comparison isolates accumulation precision while keeping precision layouts and mixing ratios matched.
  • G.4 Larger-Model Accumulation Stability: FP16 accumulation remains close to FP32 accumulation on LLaMA 3.1 8B and Qwen 2.5 14B, with larger-model differences stable across layers.On LLaMA 3.1 8B, maximum differences are on the order of 10^-2 and most means around 10^-3 or lower; on Qwen 2.5 14B, means remain around 10^-4.
  • G.5 Positional Routing and Heavy-Hitter Exposure: SpTrans25 routes only about 8.5% of selected high-importance attention mass to INT8, below its nominal 25% tile-group coverage.Lower weighted INT8 exposure indicates stronger protection of high-importance attention mass from INT8 routing.
  • L Max Mean | L Max Mean | L Max Mean | L Max Mean: Across 25%, 50%, and 75% coverage levels, structured spatial routing retains more high-importance interactions in FP16 without online detection.The exposure metric quantifies the fraction of selected important query-key mass routed to INT8.
Loading 2608.17336v1…