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

📄 arXiv: 2608.17336v1 📥 PDF

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

分类: cs.AI

发布日期: 2026-08-18

🔗 代码/项目: GITHUB

💡 一句话要点

提出TileMix以解决大语言模型推理中的计算与内存瓶颈问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长上下文 大型语言模型 推理加速 精度路由 计算效率 内存优化 自注意力机制

📋 核心要点

  1. 长上下文预填充在大型语言模型中导致计算和内存流量显著增加，现有方法未能有效解决这一问题。
  2. TileMix通过将注意力矩阵划分为硬件对齐的分数块，实现了精度的空间路由决策，优化了计算效率。
  3. 在多个基准测试中，TileMix显著提高了预填充吞吐量，并恢复了长上下文的质量，展示了其优越性能。

📝 摘要（中文）

长上下文预填充在大型语言模型（LLMs）中会导致显著的计算和内存流量，因为密集自注意力计算的查询-键分数呈二次方增长。现有方法要么采用统一的低精度路径，要么选择令牌交互，未能在融合的密集注意力中实现空间精度路由。我们提出了TileMix，这是一种以块为中心的精度路由内核，使数值精度成为在融合密集注意力中对分数块组的可执行空间决策。TileMix将注意力矩阵划分为硬件对齐的分数块，打包路由决策为紧凑的位掩码，并通过FP16或INT8分数计算调度每个块组，同时两个路径更新共享的在线软最大状态。在长上下文中，TileMix通过路由所有合法的块组，保持了密集的令牌连接性，无需训练，并支持分组查询注意力、可变长度批次和INT8键/值缓存。在LLaMA、Qwen和Vicuna的LongEval、LV-Eval和A100预填充基准测试中，TileMix恢复了在统一INT8下丢失的长上下文质量，并提高了FP16的预填充吞吐量，实现了跨模型系列的可控准确性-效率边界。

🔬 方法详解

问题定义：论文要解决长上下文预填充在大型语言模型中导致的计算和内存流量问题。现有方法要么采用统一低精度，要么选择性地进行令牌交互，未能有效利用硬件资源。

核心思路：TileMix的核心思路是将注意力矩阵划分为硬件对齐的分数块，并通过紧凑的位掩码实现精度的空间路由决策，从而优化计算效率和内存使用。

技术框架：TileMix的整体架构包括分块的注意力矩阵、路由决策的位掩码、FP16和INT8的分数计算路径，以及共享的在线软最大状态更新模块。

关键创新：TileMix的主要创新在于其块中心的精度路由机制，使得数值精度成为可执行的空间决策，保持了密集令牌的连接性，并支持多种查询方式。

关键设计：TileMix设计了紧凑的位掩码来打包路由决策，允许每个路由位控制多个相邻的键块，同时保持硬件对齐的计算块和紧凑的元数据。

🖼️ 关键图片

fig_0
img_1
img_2

📊 实验亮点

在LLaMA、Qwen和Vicuna的多个基准测试中，TileMix在统一INT8下恢复了长上下文的质量，并在FP16的预填充吞吐量上实现了显著提升，展示了其在不同模型系列中的可控准确性与效率的平衡。

🎯 应用场景

TileMix的研究成果在大型语言模型的推理加速中具有广泛的应用潜力，尤其是在需要处理长上下文的自然语言处理任务中。其优化的计算效率和内存使用将推动更高效的模型部署，适用于实时应用和资源受限的环境。

📄 摘要（原文）

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.