DEV Community

Papers Mache
Papers Mache

Posted on

Mixed‑precision routing speeds up attention prefill

TileMix doubles prefill throughput while leaving the model untouched. The speedup comes from routing attention‑score tiles to INT8 Tensor Cores instead of running everything in FP16. By keeping the dense connectivity graph intact, the method avoids any retraining or architectural tweaks. “TileMix improves throughput over FlashAttention by routing score-tile groups to INT8 Tensor Cores.” [1]

Before TileMix, practitioners had to choose between a uniform low‑precision path that hurt long‑context quality or token‑selection schemes that broke the full attention graph. FlashAttention remains the de‑facto FP16 fused kernel for dense self‑attention, and existing mixed‑precision tricks apply a single precision globally rather than per tile. Consequently, speed gains were modest and often came at the cost of accuracy on long prompts.

According to the paper, SpTrans75 reaches 31.80 K tokens /s on a 4 k‑token prompt, more than double FlashAttention’s 14.33 K tokens /s while staying ahead of other INT8 baselines. The paper does not provide a direct TileMix vs. FlashAttention throughput comparison for this setting. [1] The result holds across LLaMA, Qwen and Vicuna models, showing a consistent accuracy‑efficiency frontier.

TileMix’s design hinges on hardware‑aligned score tiles and compact bitmask routing, which may limit portability to GPUs lacking INT8 Tensor Core support or to future attention kernels that change tile granularity. Although the authors emphasize that “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,” [1] it remains an open question how well the approach scales to non‑A100 architectures or to inference frameworks that cannot fuse the mixed‑precision path.

Deployments that spend most of their latency in prefill should replace FlashAttention with TileMix on supported hardware, cutting end‑to‑end latency roughly by half without sacrificing model quality.

References

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

Top comments (0)