metal-flash-attention

收录理由

这个项目把 FlashAttention 算法移植到 Apple 芯片的 Metal 上，代码本身和工程笔记都值得细读。实现保持精简，是一组 Swift 源码，运行时 JIT 编译，而不是预编译的二进制。README 把各种硬性取舍讲得很透。Apple 硬件没有原生 FP32 原子操作，加上寄存器压力，迫使作者重新设计了反向传播：拆成独立的 dQ 和 dK/dV 内核，用更多 GEMM 换来了注意力矩阵两个维度上的完全并行化。性能数据用每秒十亿指令数（GI/s）对照 roofline 模型给出，而不是裸的 GFLOPS——在 M1 Max 上稳定跑到 4400 GI/s，ALU 利用率 83%。这个数字给了你一个具体参照，可以判断自己的 Metal shader 是否浪费了 ALU 吞吐。寄存器溢出和分块策略，对任何想把 CUDA 注意力内核搬到 Apple GPU 的人，都有直接借鉴意义。