llama.cpp b10785发布:Metal后端新增稀疏Flash Attention,prefill阶段默认启用
该版本在Metal后端为Flash Attention加入稀疏掩码支持,通过索引压缩与单遍扫描降低掩码内存访问,并将稀疏路径用于prefill。
llama.cpp发布b10785版本,在Metal后端为Flash Attention加入了稀疏掩码支持,并将该路径用于prefill阶段。
新增的kernel_flash_attn_ext_vec_idx内核会把有限的掩码条目压缩成每行索引列表(使用Hillis-Steele扫描,每个线程组处理一行),随后Flash Attention向量内核可选择性地按索引gather(增加FC slot 5)。
稀疏路径的启用条件
主机端加入门控:当n_kv_max > 0、掩码存在,且head大小与KV类型受支持、n_kv_max为 1、支持nr23变体、sinks、ALiBi、softcap、permute、v_view_of_k时走稀疏路径;不满足时回退到无掩码路径。
修复与优化
修复了稀疏Flash Attention的行寻址问题:掩码参数类型为half* 但步长按字节计算,之前每行掩码偏移被放大了 2 倍,现已改为先转为char* 再应用字节步长;稀疏pidx参数为char*,之前每行元素偏移被低估,现已按sizeof(int) 缩放以得到正确字节偏移。这两项修复解决了多行(nb*nr23[1] > 1)场景下的失败。
- 索引内核原先需要读取掩码行两次:一次计数有限条目(用于前缀扫描),再一次恢复位置。由于该内核受内存带宽限制,这使掩码流量翻倍。新实现改为单遍:在计数阶段将有限位置保留在每线程寄存器数组并直接写出,避免第二次掩码读取。
测试与配套变更
增加了稀疏Flash Attention prefill的性能测试用例,可测不同KV大小、n_kv_max提示和批大小;运行方式为 ./build/bin/test-backend-ops -b MTL0 -o FLASH_ATTN_EXT -p "n_kv_max=[1-9]" perf。
代码提交中曾启用Qwen4并调整nsg,但后续版本已禁用Qwen4。
b10785提供了macOS(Apple Silicon、Intel)、Linux(多种后端)、Android、Windows(CPU、CUDA、Vulkan等)以及XCFramework和UI的预编译二进制。