TileMix：给AI大模型的"注意力"做一次精准的精度分配手术

你有没有想过，当你让ChatGPT这类大语言模型读一篇几万字的长文档时，它到底在"卡"哪里？

答案往往不是它读不懂，而是它算不过来。模型在读长文本的第一步，也就是所谓的"预填充"阶段，需要让文本里的每一个词和其他所有词都两两比对一次，计算它们之间的关联程度。这个计算量会随着文本长度的平方增长，文本翻一倍，计算量涨四倍。这就是为什么处理长文档时，模型响应会明显变慢。

这个环节有个专业名字。

> 自注意力机制：大语言模型理解上下文的核心方法，让每个词都能"看到"并衡量与其他所有词的关联程度，从而理解语义。

面对这个平方级增长的计算量，过去几年研究者们主要想了两条路子。一条是降低数字精度，比如把原本用16位小数表示的数值，压缩成8位整数来算，这样速度快、省内存，但精度会打折扣。另一条是减少计算量本身，干脆不去算某些词之间的关联，用稀疏的连接模式代替完整的两两比对。

这两条路子都有效，但都留了一个空白地带没人碰过。这篇来自北德克萨斯大学和圣路易斯大学团队的论文，叫做TileMix，就是盯着这块空白下的手。

先搞懂：这块"空白地带"到底是什么

要理解TileMix填的是什么坑，得先弄明白现在主流的高效注意力计算是怎么运作的。

现在业界公认最能打的注意力实现方案叫FlashAttention，它的核心思路是把一个巨大的注意力计算矩阵切成很多小方块（术语叫"tile"，瓦片），一块一块地搬进GPU的高速缓存里算，算完就扔，不把整个大矩阵都摊在显存里占地方。这样既省显存又提速。

> FlashAttention：一种通过分块计算和流式处理注意力矩阵的技术，避免把巨大的中间结果直接摊在显存里，从而节省内存带宽并提速。这是当前几乎所有高效注意力实现的技术基础。

这个方案有个特点：不管切成多少块，每一块内部用什么精度去算，是固定死的。要么整个注意力计算过程都用16位浮点数（FP16），要么统一换成8位整数（INT8），没有中间地带。

而稀疏注意力那条路呢，是在决定哪些词对之间要不要建立连接，跟精度完全是两码事。它管的是"连不连"，不管"怎么算"。

TileMix团队的洞察是：这两条路子分别管住了"哪里计算"和"用什么数值格式计算"这两个维度里的一个，但从来没有人把"每一小块矩阵该用什么精度"当成一个可以按空间位置单独决定的执行选项。换句话说，之前的方法要么是全局统一精度，要么是全局统一连接方式，但没人试过让精度本身像连接方式一样，按照矩阵里的位置来分区管理。

这就好比一家餐厅的后厨，以前要么整个后厨统一用高压锅（快但可能把食材压得不够精细），要么整个后厨统一用小火慢炖（精细但慢），从没人想过按菜品所在的灶台分别配置：靠近门口的几个灶台用高压锅处理不太讲究的家常菜，最里面几个灶台用小火慢炖伺候需要精细火候的招牌菜。如果不这样区分，你要么牺牲了所有菜品的口感去换速度，要么牺牲了所有出餐速度去保口感，没有折中的空间。

TileMix做的事情，就是把这种"按灶台分配烹饪方式"的思路，搬进了注意力矩阵的计算里。

TileMix到底怎么设计的

### 第一步：把矩阵切成方块，再给每块贴标签

TileMix的基本操作单元，延续了FlashAttention的思路，把整个注意力矩阵切成一个个硬件友好的小方块。

> 硬件对齐的计算块（tile）：为了配合GPU的计算单元（Tensor Core）效率最高的处理粒度而设计的固定尺寸矩阵分块，通常是几十到一百多行列的小方阵。

但TileMix多做了一件事：给每一小块（或者说每一组相邻的小块）贴上一个二进制标签，标签只有两种取值，0代表这一块用FP16去算，1代表用INT8去算。

这个标签的选取不是随机的，也不需要额外训练模型去学习，而是用几种事先设计好的固定模板来分配。这些模板本身借鉴了之前稀疏注意力研究里总结出的几种典型注意力分布规律，比如局部窗口模式、全局关键位置模式、随机采样模式等等，只不过之前这些模式是用来决定"要不要计算"，现在被TileMix拿来决定"用什么精度计算"。

这里有一个关键区别必须说清楚：TileMix从头到尾都保留了全部的词对连接，一个都没有砍掉。它只是在"连接保留"的前提下，决定每个连接用什么精度去计算。这跟稀疏注意力那种直接切掉部分连接的做法有本质不同。

### 第二步：把决策压缩进一个64位的数字里

一个长文本可能要切成几十甚至上百个小方块，如果给每个块单独存一个标记位，管理起来会很麻烦，查找的时候也慢。TileMix的解决办法是把一整行（对应一个查询块）里所有相关方块的精度决策，打包压缩进一个64位的整数里，每一个二进制位对应一组方块的精度选择。

> 位掩码（bitmask）：把多个二进制判断结果压缩存储在一个整数的不同二进制位上的技术，查询时只需要做一次位移和按位与运算就能取出结果，速度极快。

想取出某一组方块的精度决策时，只需要做一次位移操作加一次按位与运算，这在计算机里几乎是最快的操作之一，几乎不占用时间。这就好比你要记住一栋楼里100个房间哪些开灯哪些关灯，与其为每个房间单独写一张纸条再逐张翻找，不如把这100个房间的开关状态编码成一串二进制数字存在手机备忘录里，要查第37个房间的状态，直接算出它在这串数字的第几位，一眼就能读出来。如果不这样压缩，管理这些海量的精度决策标记会占用大量额外的存储和查询开销，反而抵消了精度优化本身带来的加速收益。

当文本特别长，方块特别多，超过64个位不够用怎么办？TileMix设计了一个"分组"机制，让一个二进制位可以同时管理好几个相邻的方块，就像一个开关同时控制好几盏灯，这样即使方块数量暴涨，64位仍然够用。

### 第三步：两条路径殊途同归，共享同一套"记账系统"

这是整篇论文技术含量最高的部分。

在FlashAttention的流式计算里，每处理完一小块矩阵，都要更新一套共享的状态变量，包括当前为止见过的最大分数值、归一化分母、以及累加的输出结果。这套状态变量一直被后续的每一块计算复用，是整个流式计算能够正确工作的核心机制。

> 在线softmax：一种边流式处理矩阵块边动态更新归一化统计量的算法，不需要把整个矩阵都算完才能做归一化，这也是FlashAttention能够省内存的关键技巧。

问题在于，FP16路径和INT8路径产生分数的方式完全不同。INT8那边要先把浮点数压缩成整数，用整数乘法算完之后再乘回一个缩放系数还原成浮点数，这个过程会引入舍入误差和精度损失，跟FP16直接浮点运算的行为不一样。

TileMix的做法是：不管这一块是走FP16还是INT8算出来的，在进入共享的状态更新之前，都要先统一转换到同一个浮点数值域里去。这样两条路径产生的分数值，虽然来源不同，但最终都以相同的"语言"汇入同一套记账系统，不会因为混用两种精度而搞乱整体的归一化过程。

这就像一家公司同时收到人民币账单和美元账单，财务在做总账之前，必须先把美元按当天汇率换算成人民币，再放进同一张总账表里加总。如果不做这一步统一换算，直接把两种货币的数字硬凑到一起相加，得出来的总数毫无意义。TileMix做的正是这件"汇率换算"的工作，只不过换算的对象是数值精度而不是货币。

支持的场景：不只是理论上能跑

论文里特别强调了TileMix在实际部署场景里的几个支持能力。

它支持分组查询注意力，这是现在主流大模型（包括Llama、Qwen这些）普遍采用的一种省显存的注意力变体，多个查询头共享同一组键值头。TileMix让共享同一个键值头的查询头，也共享同一套精度路由决策，不需要重复计算。

它支持变长批处理，也就是同一批次里不同长度的文本可以混着一起处理，不需要都填充到同样长度浪费计算资源。

它还支持INT8格式的键值缓存，这个跟生成阶段（而不是本文重点讨论的预填充阶段）的效率有关，缓存里存的键值对用低精度存储可以省下大量显存。

这几项支持能力凑在一起，意味着TileMix不是一个只能在实验室跑通的原型，而是真的考虑了实际部署会遇到的各种边角场景。

实验说了什么：精度分配确实能省下时间，还能保住准确率

论文在两个长文本理解任务上做了测试。

第一个任务叫LongEval，说白了就是让模型在一堆几万字的文本里，精确找出某一行特定编号对应的内容，看它能不能一字不差地把答案抠出来。

第二个任务叫LV-Eval，覆盖了11个中英文长文本问答数据集，涉及事实核查、多跳推理、多字段问答等各种题型，文本长度从16000字到64000字不等。

结果显示，把注意力全部换成INT8（论文里叫One配置）之后，模型的表现普遍会明显下降。以Llama 3.2 3B模型为例，在64000字长度的多个数据集上，纯INT8配置的得分经常只有FP16原始配置的六到七成，掉分幅度相当可观。

但如果引入TileMix的精度混合路由，情况就不一样了。比如在16000字长度的"事实召回"这个子任务上，一种叫SpTrans（借鉴了稀疏Transformer设计思路的精度布局）的配置，得分能达到21分左右，而对应的纯FP16基线只有6.72分，纯INT8配置更是只有4.45分。这个结果在Qwen 2 7B和Qwen 2.5 7B模型上同样复现出来了，说明这不是偶然的模型特异性现象，而是一种跨模型的规律性发现。

效率方面，在4000字长度的Llama 3.2 3B模型测试里，TileMix的某种混合配置吞吐量达到每秒31.8万token，而标准FlashAttention只有每秒14.33万token，纯INT8配置是每秒29.8万token。也就是说，TileMix的混合路由方案不仅没有比纯INT8慢，反而在这个测试点上跑得更快，同时还保留了远比纯INT8丰富的模型质量。

下面这张表格摘录了论文中一部分关键的效率对比数据（单位为每秒千token数）：

| 文本长度 | Torch标准实现 | FlashAttention | 纯INT8(One) | TileMix混合(75%INT8) |

| 1k | 11.14 | 17.45 | 32.27 | **33.50** |

| 2k | 7.78 | 16.48 | 32.06 | **33.92** |

| 4k | 显存溢出 | 14.33 | 29.80 | **31.80** |

| 8k | 显存溢出 | 显存溢出 | 27.41 | 26.61 |

可以看到，随着文本变长，标准的Torch实现和普通FlashAttention都相继出现显存溢出问题，只有做过精细内存管理的TileMix和纯INT8方案还能撑住，而TileMix在多数长度上还能进一步领先纯INT8。

数字精度混用到底靠不靠谱：论文做的"体检"

光看任务表现还不够，研究者还专门做了一轮数值层面的"体检"，检验FP16和INT8混合计算之后，模型输出到底跟纯FP16参考结果差多少。

结果显示，INT8覆盖比例从0%涨到25%的过程中，输出数值的平均绝对偏差是逐渐增大的，但这个增大过程是可控、可预测的，覆盖率越低偏差越小。这说明INT8覆盖率本身可以当作一个实用的"数值控制旋钮"，需要更保守就调低覆盖率，能接受更多误差就调高覆盖率换取更快速度。

这套体检还测试了不同模型深度的影响，发现层数越多的模型（比如32层的模型相比12层模型），数值偏差累积得更明显，这符合直觉，因为误差会一层一层往下传递、放大。

还有一个挺有意思的补充实验，检验的是"静态路由模板到底有没有把宝贵的高精度资源用在刀刃上"。研究者专门计算了注意力矩阵里哪些位置的关联程度对最终结果影响最大（称为"重要交互"），然后看SpTrans这种精度布局模板有没有把这些重要交互优先分配给FP16高精度路径。结果显示，在名义上25%的方块被分配给INT8的情况下，实际被划入INT8的高重要度交互只占到了8.57%，远低于25%这个整体比例。换句话说，这套静态模板虽然没有用到任何实时的重要性检测机制，却已经天然地把更多计算资源留给了那些真正举足轻重的词对关系。

这就好比一个仓库管理员，即便完全不知道具体哪些货物是紧俏商品，只是凭着"靠近出入口的货架通常放常用品，最里面角落放冷门货"这条经验规则去摆放货物，结果却往往歪打正着地把真正的畅销品摆在了顺手可取的位置。这说明经过大量真实场景验证提炼出的空间分布规律，本身就自带了某种"重要性感知"的能力，即使没有专门为此设计检测机制。

写在后面

读到这篇论文时，最让我意外的不是它提出了一个新点子，而是它指出的那个"空白地带"竟然一直没人碰过。低精度量化和稀疏注意力这两条路线各自发展了好几年，居然从来没有人想过把精度当成一个可以按空间位置分配的资源。这有点像两拨人各自在自己的赛道上狂奔，谁都没想过赛道之间其实可以架一座桥。

论文里那个"重要交互暴露率"的实验，是我觉得最值得单独拎出来说的细节。研究者没有满足于"我的方法效果好"，而是反过来去问"为什么静态模板会有效"，结果发现这些借鉴自稀疏注意力研究的空间分布规律，本身就带着某种朴素的重要性直觉。这提醒我一件事：很多领域里积累下来的经验模式，即便原本是为了解决A问题设计的，挪到B问题上可能依然管用，因为它们背后反映的是数据本身的某种结构性规律，而不是针对具体任务的死板技巧。

这篇论文目前只在英伟达A100这一款GPU上做了验证，也只测试了FP16和INT8这一对精度组合。如果换成更新的GPU架构，或者尝试FP8、INT4这些精度格式的组合，这套精度路由的思路是不是依然有效，会不会需要重新设计位掩码的编码方式，这是个留白的问题，也是个值得继续追问下去的方向。

Q&A

Q1：TileMix是什么，它解决了什么问题？

A：TileMix是一种针对大语言模型长文本处理的注意力计算加速方法，它把注意力矩阵切成小方块，让每个方块可以单独选择用FP16高精度还是INT8低精度计算，同时保留全部词对连接不做删减，从而在保持模型理解质量的同时提升长文本预填充阶段的计算速度。

Q2：TileMix和普通的INT8量化有什么区别？

A：普通INT8量化通常是整个计算过程统一换成低精度，而TileMix允许在同一次注意力计算里，不同的矩阵区块分别选择FP16或INT8，通过打包的位掩码实现按区域灵活分配精度，实验显示这样比统一使用INT8能更好地保住长文本理解的准确率。

Q3：使用TileMix之后模型速度能提升多少？

A：论文实验显示，在4000字文本长度下，TileMix某些混合精度配置的吞吐量能达到每秒31.8万token，高于标准FlashAttention的每秒14.33万token，也略高于纯INT8方案的每秒29.8万token，同时在长文本问答准确率上明显优于纯INT8方案。