TileMix Routes Attention Precision at the Hardware-Tile Level

Introduction

Long-context inference places a particularly heavy burden on the prefill stage. Before generation begins, the model must process a large input sequence and construct attention results across many query-key pairs. Dense self-attention therefore faces both rapidly growing computation and substantial memory traffic. Lowering precision can help, but applying INT8 uniformly may recover efficiency at the cost of long-context quality.

TileMix addresses this tension by asking a more specific question: which regions of the attention score matrix need higher numerical precision?

Precision routing over score tiles

The method partitions the attention matrix into score tiles aligned with the way accelerator kernels execute work. Neighboring key tiles can be grouped together, and a compact bitmask records whether each group should use FP16 or INT8 score computation. The selected groups are then dispatched within one fused, FlashAttention-style kernel rather than being handled by separate attention operators.

The grouping mechanism is important for practical execution. A decision at every individual interaction would create excessive metadata and control overhead, while a decision at an overly large granularity would limit the ability to preserve precision where it matters. TileMix allows one routing bit to govern multiple adjacent key tiles. This keeps the underlying compute tiles hardware-friendly and makes the routing representation more compact as context length grows.

FP16 and INT8 groups also share the same online-softmax state. Each path contributes to a common running normalization process, allowing mixed-precision tiles to participate in one dense attention computation. The design is consequently more than a static quantization recipe: precision becomes a spatial decision executed inside the attention kernel.

Dense connectivity instead of token pruning

TileMix does not decide whether a token interaction exists. It routes legal tile groups to different numerical paths, while retaining dense token connectivity. That distinction separates it from approaches that reduce cost by selecting or pruning interactions. The paper presents the method as training-free and compatible with grouped-query attention, variable-length batches, and INT8 key/value caches.

What the reported results suggest

The authors evaluate TileMix on LongEval, LV-Eval, and A100 prefill benchmarks across LLaMA, Qwen, and Vicuna model families. According to the supplied paper summary, the method recovers some of the long-context quality lost under uniform INT8 and improves prefill throughput relative to FP16. Different routing choices can therefore expose a controllable accuracy-efficiency frontier instead of forcing a binary choice between full FP16 and uniform low precision.

The broader significance is an implementation-level reframing of mixed precision. TileMix does not remove the quadratic nature of dense attention, nor does it claim to replace sparsity with a universal solution. Its contribution is to align precision decisions with the spatial structure already used by accelerator kernels and fused attention. The practical outcome will still depend on routing policy, sequence length, hardware, and model family, making the released implementation and detailed benchmark settings important for deployment assessment.

Source: Hugging Face Daily Papers