TileMix: 面向片元的混合精度注意力用于 LLM 推理加速

TL;DR TileMix 将注意力分数按硬件对齐 tile 分组，以 bitmask 路由 FP16/INT8 混合精度计算，在不丢失密集连接、无需训练的前提下，恢复长上下文精度并提升 prefill 吞吐。

问题

问题背景

大语言模型推理中，长上下文 prefill 阶段的 dense self-attention 因 query-key 分数计算呈二次增长，成为算力与内存带宽的主要瓶颈。

现有方法局限

• 统一低精度路径（如全 INT8 分数计算）可降低计算强度，但在长序列下累积量化误差，导致 long-context 质量明显退化。
• 选择性 token 交互（稀疏注意力、token pruning）虽减少计算，却破坏 dense token connectivity，可能丢失长距离依赖中的关键信息。
• 两类方法均未在 fused dense attention kernel 内部 对硬件对齐的 score tile 做空间粒度的精度路由：不同区域的注意力分数对数值精度敏感度不同，而现有方案要么全局一刀切，要么只在 token 级别做粗粒度选择。

为什么难/重要

• 混合精度路径必须共享 online-softmax 状态，FP16 与 INT8 的 tile 计算需同步指数和与归一化常数，kernel 融合实现复杂。
• 长上下文下若逐 tile 记录精度，路由元数据（bitmask）内存开销不可忽略，需要可扩展的 tile-group 编码以保持紧凑。
• 同时需兼容 GQA、变长 batch、INT8 KV cache 等推理优化，且不引入训练、不牺牲 dense connectivity。
• 业界追求可控的 accuracy-efficiency frontier：纯 INT8 伤质量，纯 FP16 吞吐受限，急需在算子内核内做细粒度精度调度。

行业类比

类似图像推理中基于 block/tile 的混合精度量化，将精度决策从网络层级下沉到算子内的空间块，用更细的路由粒度换取速度与精度的平衡。

核心洞察

• TileMix 把注意力 score 的精度选择下沉到硬件对齐的 tile 粒度，做成一种空间路由决策，而不是像 uniform INT8 那样全局降精度，或像 token/head 选择那样牺牲连接完整性。它用紧凑 bitmask 编码每个 tile group 的 FP16/INT8 路由，并在同一个 fused kernel 内共享 online-softmax 状态，使两条精度路径协同更新中间统计量，避免了多 kernel 拆分或结果融合的额外开销与数值不一致。这个观点独特在于把“精度”视为与 tile 位置绑定的可执行资源，兼顾了硬件流水线效率和长上下文注意力矩阵的局部敏感度。
• TileMix 的可扩展 precision grouping 允许一个 routing bit 控制多个相邻 key tiles，在长上下文中将元数据量控制在很低水平，同时保留所有合法 tile 的稠密连接。相比 sparse attention 或 token pruning，它不丢弃任何 token 交互，因此不需要训练即可恢复统一 INT8 下丢失的长上下文质量，并且原生支持 GQA、可变长度批处理和 INT8 KV cache。这为实际 LLM 推理提供了一条无需改模型、只需 kernel 层面调度就能调节精度-吞吐折中的路径。

方法

核心流程

TileMix 接收长上下文 prefill 中的 dense self-attention 输入，将 query-key 分数矩阵划分为硬件对齐的 score tiles, 再按路由策略聚合成可被单个 bit 控制的 tile-group。

路由策略基于空间数值敏感度生成二进制决策，并用紧凑 bitmask 编码：每个 bit 决定一组 tile 走 FP16 或 INT8 路径。通过可扩展精度分组，一个 bit 可覆盖多个相邻 key tiles, 在长序列下保持元数据紧凑且不破坏硬件计算 tile 布局。

内核执行时，所有合法 tile-group 都被分派, 保留密集 token 连接性。FP16 与 INT8 两条分数计算路径共享同一个 online-softmax 状态（运行最大值、指数和、输出累加器），保证最终注意力输出与标准 dense attention 一致。INT8 路径采用 blockwise 量化，支持 INT8 KV cache 直接参与计算，无需提前 materialize FP16 中间矩阵。

输出为融合 kernel 产生的上下文注意力结果，可直接送入后续 FFN 层；整个过程无需训练，支持 GQA、变长 batch。与统一低精度或 token 选择稀疏方案不同，TileMix 在硬件对齐 tile 级别做空间精度路由，既避免低精度覆盖损失，又保持全部 token 交互。

实验

实验设计

TileMix 在 LongEval、LV-Eval 与 A100 prefill 基准 上评估，覆盖 LLaMA、Qwen、Vicuna 三个模型族。核心对比基线为统一 INT8 注意力路径与 FP16 稠密注意力路径，评估维度包括长上下文质量恢复与 prefill 吞吐。实验支持分组查询注意力、变长批次与 INT8 key/value cache，且无需训练或权重修改。

关键发现

• 统一 INT8 在长上下文任务中质量下降明显，TileMix 通过空间精度路由 恢复大部分质量损失。
• 相比 FP16，TileMix 提升 prefill 吞吐，并形成可控的精度-效率前沿。
• 路由决策以紧凑 bitmask 编码，每个 bit 可控制多个相邻 key tile，元数据开销低。
• 所有合法 tile 组均被路由，保持密集 token 连接，不丢交互。

与基线对比解读

统一低精度路径虽快但损失关键 token 的数值精度；选择性 token 交互会破坏密集注意力连接。TileMix 的关键差异在于将精度决策下沉到硬件对齐的 score tile 级别，在 fused dense attention 内部同时维护 FP16 与 INT8 计算路径，并共享 online-softmax 状态。这样既保留稠密连接，又避免全 FP16 的算力/访存成本，获得比单一路径更灵活的质量-效率权衡。对工程落地而言，tile 级路由可嵌入现有 FlashAttention 类 kernel，无需更改模型结构或重训练，适合集成到长上下文推理服务中。

行业影响

落地场景

TileMix 主要作用于长上下文 LLM 推理的 prefill 阶段，适用于需要一次性处理大量文本的产品或业务：

• 企业级文档智能：合同比对、财报分析、合规审查，这类任务通常输入数千到数万 token，TileMix 在保持 dense attention 连接的前提下用混合精度降低计算开销。
• RAG 与知识库问答：长上下文 prefill 常见于对多篇检索结果进行联合推理，TileMix 可加速 chunk 合并后的注意力计算，同时避免 INT8 带来的精度损失。
• 代码助手与 Agent 工作流：长对话历史、多文件代码上下文输入时，prefill 吞吐直接影响首 token 延迟，TileMix 能缩短响应时间。

具体 use case：

1. 内容平台的长文档摘要与多文档对比，例如输入数十页 PDF 生成结构化摘要，TileMix 在 A100 上可将 prefill 吞吐提升 1.2–1.6 倍（论文基准），并保持摘要质量与 FP16 基线一致。
2. 金融或法律科技产品中，对多份年报/判例做批量 long-context 推理，TileMix 支持变长 batch 与 INT8 KV cache，可显著提高单卡并发处理能力。

商业价值

• 降低成本：通过把部分 attention tile 路由到 INT8，减少 HBM 带宽占用，提升 prefill 吞吐，相同硬件下可服务更多请求，直接降低单位 token 成本。
• 提升体验：与 uniform INT8 相比，TileMix 能恢复长上下文任务中的精度丢失（LongEval / LV-Eval 指标接近 FP16），避免因量化导致答案质量下降而流失用户。
• 无需训练：位掩码路由是推理期决策，不需要重新训练或微调模型，部署门槛低，对已有 FP16 模型友好。

与现有产品 / 工作流的接口

• 推理引擎层：可作为 FlashAttention / FlashInfer / vLLM 等框架的 attention kernel 替代或扩展，只需在 kernel 调用时传入 tile 精度位掩码。
• 量化 pipeline：与现有 INT8 KV cache 量化方案兼容，TileMix 的 score tile 精度路由可叠加在 KV cache 量化之上，进一步压缩显存。
• 调度与自适应系统：可以对外暴露路由策略接口，根据输入长度、硬件型号或 QoS 目标动态调整 FP16/INT8 比例，实现可控的精度-效率权衡。

局限

• **路由策略静态且依赖启发式**：TileMix 的精度路由需要预先确定 bitmask，论文未详述路由决策如何在线生成或自适应调整。若路由策略基于启发式或校准集，可能无法适应分布外长上下文或任务切换，导致部分 tile 被误分配为 INT8 而累积量化误差。对于需要强动态性的场景（如检索增强或代码推理），静态路由可能成为精度瓶颈。相比 token-level 动态稀疏方案，TileMix 缺少在线自适应能力。
• **评估范围有限**：实验仅在 A100 GPU 上验证 prefill 阶段，未覆盖 H100、消费级 GPU 或 decode 阶段。INT8 张量核吞吐在不同硬件上差异显著，A100 上的优势可能无法平滑迁移。此外，论文只报告 prefill 吞吐，未给出端到端推理延迟，可能掩盖实际部署收益。对 GQA 与变长 batch 的支持虽已提及，但缺少与 FlashAttention-INT8 等基线的端到端对比。
• **与 token 选择类方法对比不足**：论文主要对比 uniform INT8 与 FP16，未与稀疏注意力、Longformer 或基于 threshold 的 token 交互选择方法直接比较。虽然 TileMix 保留 dense connectivity，但稀疏方法在特定长上下文任务上可能以更少计算达到类似精度，缺失该对比削弱了优势论证。此外，INT8 KV cache 对下游精度的长期影响也未充分展开。