深入 Android 端侧 AI 推理的 Transformer 算子优化全链路：从 Self-Attention 计算瓶颈到移动 GPU 计算着色器的 FlashAttention 加速
在端侧跑 1B 参数 LLM，prefill 阶段延迟总卡在 attention 算子。用 Mali Offline Compiler 和系统 profiler 抓数据，发现瓶颈不在 GEMM 算力，而在中间张量的读写。本文记录我把 FlashAttention 搬进 OpenGL ES Compute Shader 的过程。
Self-Attention 的瓶颈在访存不在算力
标准 attention 流程：S = QK^T，P = softmax(S)，O = PV。
序列长度 n 时，中间矩阵 S 是 n×n。n=2048、fp16 下 S 占 8MB。这块矩阵要先写回全局内存，再读出来做 softmax，再读出来乘 V，三次全量读写。
移动 GPU 的 LPDDR 带宽比桌面显卡差一个数量级，进一步放大了这个矛盾。我在 Mali-G710 上实测，attention 算子的算术强度只有 GEMM 的 1/20 左右，计算单元大部分时间在等数据。
FlashAttention 的本质：重排计算顺序
FlashAttention 不改变数学结果，只改变计算顺序，核心就两点：
- 把 Q、K、V 切成小块，每次只算一个小 tile 的 attention，不生成完整的 S 矩阵。
- 用 Online Softmax 增量维护最大值和分母，最后统一归一化。
分块带来的是数值稳定性问题：softmax 减最大值是为了防指数溢出，但分块处理时新块的最大值可能更大，需要把之前累计的结果按比例缩放回去。
# Online Softmax 分块更新
m_new = max(m_old, max(S_block)) # 更新 running max
scale = exp(m_old - m_new) # 旧累计的缩放因子
l_new = l_old * scale + sum(exp(S_block - m_new))
O_new = O_old * scale + P_block @ V_block # 加权累加输出这样每个 tile 只进出内存一次，全局访存从 O(n²) 降到 O(n)。
移动 GPU 的约束：Compute Shader 怎么落地
端侧实现我选了 OpenGL ES 3.1 的 Compute Shader，而不是 Vulkan。理由很直接：老设备覆盖率高，NNAPI 的 GPU 后端也走这条技术路线，后续可以复用。
移动 GPU 和桌面 GPU 两个差异直接影响实现：
- 没有高带宽 L2，shared memory 很小。Mali 的 workgroup 内存通常在 16KB 量级，tile 大小要贴着这个上限定。
- 工作组分发开销高，所以一个 workgroup 要一次算完整一个输出 tile，内部循环把 K 方向的所有分块都迭代完。
GLSL 的 shared 内存声明：
// 每个 workgroup 缓存一个 K 方向 tile
shared mediump float Qs[BR][BD];
shared mediump float Ks[BC][BD];mediump float 在 GLSL ES 规范里只保证 10 位精度。真正想用 fp16 存 shared，得看扩展支持，很多设备不支持就直接退化成 fp32，带宽翻倍，后面踩坑就在这。
一个完整的 tile 循环
核心循环结构：外层遍历 Q 的 row block，内层遍历 K/V 的 column block。
for (int j = 0; j < TC; j++) { // 遍历 K/V 列块
load_tile(Ks, K, j); // 协作加载 K tile
float S[BR][BC];
matmul(S, Qs, Ks); // S = Qs @ Ks^T
float m_new = max(m_old, row_max(S));
float scale = exp(m_old - m_new);
l = l * scale + row_sum(exp(S - m_new));
load_tile(Vs, V, j);
O = O * scale + exp(S - m_new) * Vs; // 累加输出
m_old = m_new;
}S 全程留在寄存器里，从不写回全局内存，这是省带宽的关键。workgroup 结束时用 l 做一次归一化输出 O。
踩过的坑
fp16 精度：Mali 上 shared 内存存 fp16 需要 GL_EXT_shader_16bit_storage，不少中端芯片不支持，只能退 fp32，tile 大小直接砍半。我的做法是编译期生成两个 variant，运行时探测扩展再选路径。
bank conflict：加载 K tile 用 mat4 向量化读取能减少指令数，但 shared 内存布局不对会引入 bank conflict。让相邻线程访问相邻 bank，tile 按行主序还是列主序存储要想清楚，我用列主序后冲突明显下降。
数值稳定性：scale = exp(m_old - m_new) 在差距大时会下溢为 0，数学上正确，但部分 GPU 开启 fast-math 后会产生异常值，需要关掉 precision hint。
实测 Mali-G710 上 2048 序列的 attention 从 120ms 降到 18ms，6 倍多提升。
落地建议
选 tile 大小时，先查目标设备的 workgroup 内存上限，再留出 Q、K、V 三块加输出缓存的空间，别拍脑袋定。
编译期多生成 fp16 / fp32 两个 variant，运行时探测扩展选路径，比在 shader 里用 #ifdef 动态判断更省心。
优化顺序上先解决访存，再调 bank conflict，最后才动精度 hint。GPU 优化里 80% 的收益来自前两步，精度相关的微调经常得不偿失。