llama.cpp 项目已为 Metal 后端引入稀疏 Flash Attention (FA) 支持,专门针对模型推理的预填充阶段。此次更新包含一个新内核,将有限掩码条目压缩为每行索引列表,并通过可选的稀疏索引收集扩展了向量 FA 内核。

  • 实现中添加了一个主机端门控,在 n_kv_max 大于 0 且满足其他条件时启用稀疏路径。
  • 修复了稀疏 Flash Attention 内核中的行寻址错误,该错误此前导致多行情况下的失败。
  • 添加了优化技术以执行单次通过稀疏索引压缩,通过将有限位置保留在每个线程的寄存器中,避免冗余掩码读取。
  • 添加了测试用例和性能基准测试,以验证头大小、量化类型和各种配置参数。

当利用稀疏注意力提示时,此更改通过减少预填充阶段的内存流量来提高 Metal 用户的效率。