ik_llama.cpp PR #574 技术解析将 KQ mask 填充对齐从 16 提升到 64以适配 Vulkan coopmat2 闪存注意力【免费下载链接】ik_llama.cppllama.cpp fork with additional SOTA quants and improved performance项目地址: https://gitcode.com/GitHub_Trending/ik/ik_llama.cpp导读本文基于 ik_llama.cpp 仓库中的 PR #574《Change KQ mask padding to 64》展开。该 PR 将 Flash Attention 使用的 KQ mask 在 token 维度的填充对齐padding从 16 提升到 64以解决 NVIDIA 驱动升级到 575 后 Vulkan 后端启用 cooperative matrix 2coopmat2时触发的断言失败并同步给出了同一张 RTX 4080 上 Vulkan 与 CUDA 的 sweep-bench 性能对比。读完本文你将理解 KQ mask 填充的底层原因、它与 Vulkan coopmat2 着色器分块策略之间的约束关系以及如何复现该性能测试。一、PR 概览一次由驱动升级触发的断言修复PR #574 由 ik_llama.cpp 作者 ikawrakow 于 2025-07-03 提交并当日关闭State: Closed。PR 描述非常简短但信息量不小改动内容将 KQ mask 的填充对齐从 16 改为 64直接动机Vulkan 后端在启用 coopmat2 时需要此对齐参考基准mainline上游 llama.cpp中该值同样是 64。作者在评论区补充了触发场景他在两台远程机器中的一台将 NVIDIA 驱动更新到 575该版本开始启用 Vulkan coopmat2 能力随即在 Vulkan 后端触发了断言assert本 PR 正是为修复这一断言而提出。同时作者也借机观察了性能此前 coopmat1 路径下 Vulkan 相比 CUDA 慢了约 3 倍而根据 讨论 #562 中的预期同一张 NVIDIA GPU 上 Vulkan 与 CUDA 的差距在启用 coopmat2 后应缩小到 20%25%。遗憾的是作者在 RTX 4080 上的实测结果并未达到这一预期——coopmat2 虽有改善但 prompt processingPP仍比 CUDA 慢约 2 倍。二、背景知识什么是 KQ mask为什么要填充在自回归 Transformer 的解码与训练中注意力分数矩阵需要叠加掩码mask用来屏蔽未来的 tokencausal mask以及 padding 位置。在 ggml 中这个掩码张量被称为 KQ mask其维度布局在 ggml_flash_attn_ext 的接口注释 中有明确说明// q: [n_embd, n_batch, n_head, 1] // k: [n_embd, n_kv, n_head_kv, 1] // v: [n_embd, n_kv, n_head_kv, 1] !! not transposed !! // mask: [n_kv, n_batch_pad, 1, 1] !! n_batch_pad GGML_PAD(n_batch, GGML_KQ_MASK_PAD) !! // res: [n_embd, n_head, n_batch, 1] !! permuted !!注意 mask 的第二个维度是n_batch_pad即真实 batchtoken 数经GGML_PAD向上取整到对齐粒度后的值。填充的意义在于让 GPU 着色器可以按固定分块大小无分支地加载掩码数据如果每个工作组的行数块恰好整除填充后的维度着色器就不需要做越界 clamp从而减少边界分支、提高内存访问的规整性。这个对齐粒度正是GGML_KQ_MASK_PAD宏定义在 ggml/include/ggml.h#L2486-L2490#if GGML_USE_VULKAN #define GGML_KQ_MASK_PAD 64 #else #define GGML_KQ_MASK_PAD 16 #endif也就是说PR #574 落地后的实际状态是Vulkan 后端使用 64其余后端CPU/CUDA/Metal 等维持 16。这个条件编译结构本身就说明了问题——64 是 Vulkan coopmat2 路径特有的硬性需求而非所有后端通用的改动。三、为什么是 64coopmat2 着色器的分块约束要理解为什么必须是 64需要回到 Vulkan 后端的 Flash Attention 实现。在 ggml/src/ggml-vulkan.cpp#L2154-L2173 的fa_spec_constants中Vulkan 后端为每条 Flash Attention 管线计算工作组的分块行数rows_cols[0]并在其上方写下了关键注释与断言// mask dim1 is padded to 64, we rely on this to avoid clamping mask loads GGML_ASSERT((GGML_KQ_MASK_PAD % rows_cols[0]) 0);这条断言的含义是掩码第二维的填充量64必须能被 Flash Attention 着色器每个工作组一次处理的行数整除。若填充仍为 16而 coopmat2 着色器按更大分块64 行读取掩码就会出现 16 无法整除 64 的情况断言失败——这正是驱动升级到 575、coopmat2 能力被识别后立即触发 assert 的直接原因。同样的约束在运行时非管线构建阶段再次出现于 ggml/src/ggml-vulkan.cpp#L6434-L6435// mask dim1 is padded to 64, we rely on this to avoid clamping mask loads GGML_ASSERT((nem1 % GGML_KQ_MASK_PAD) 0);即实际推理时掩码张量的ne[1]维度必须能被GGML_KQ_MASK_PAD整除。这两处断言一前一后分别守护着着色器管线构建与执行两个阶段共同保证了掩码加载无需 clamp。从讨论 #562 中的对话可以进一步印证 64 与 coopmat2 性能的关系coopmat2 之所以在较大 KV 深度下表现明显更好正是因为它一次处理更多行64 vs 16见 讨论 #562 中关于N_KV深度影响的讨论。也就是说coopmat2 的分块粒度本身就与 64 对齐掩码填充随之对齐到 64 是顺理成章的设计。四、Coopmat2 管线驱动 575 带来的新能力Vulkan 后端的 Flash Attention 目前存在三条代码路径在 ggml/src/ggml-vulkan.cpp#L2197-L2217 中可以看到它们的创建逻辑FA_SCALAR标量路径无条件创建支持 F16、Q4_0、Q8_0FA_COOPMAT1基于VK_KHR_cooperative_matrix扩展的 cooperative matrix 1需设备报告coopmat1_fa_supportFA_COOPMAT2基于VK_NV_cooperative_matrix2扩展的 cooperative matrix 2需设备报告coopmat2支持。其中 coopmat2 路径flash_attn_cm2.comp支持的量化类型最丰富包括 F16、Q4_0、Q4_1、Q5_0、Q5_1、Q8_0 与 IQ4_NL并且额外创建了 mul_mat 的_cm2系列管线见 ggml/src/ggml-vulkan.cpp#L2221-L2229。这也是 PR 描述中所说的coopmat2 启用后需要 64 填充的完整背景coopmat2 是相对较新的 NVIDIA 扩展VK_NV_cooperative_matrix2需要较新的驱动r575 起才会被枚举出来见 讨论 #562 中关于驱动版本与matrix cores: NV_coopmat2的实测输出。因此PR #574 本质上是为新硬件能力补齐了掩码布局契约。五、对模型图构建的影响掩码张量如何被填充KQ mask 的填充并不只发生在 Vulkan 后端内部模型图的构建阶段就需要按GGML_KQ_MASK_PAD分配掩码张量。以使用 Flash Attention 的通用图构建路径 src/graphs/build_dflash.cpp#L369-L419 为例const int64_t n_kv_total GGML_PAD(ctx_len n_tokens, (int64_t) llama_kv_cache::get_padding(flash_attn)); ... lctx.dflash.inputs.kq_mask ggml_new_tensor_2d(ctx0, mask_type, n_kv_total, GGML_PAD(n_tokens, GGML_KQ_MASK_PAD)); lctx.dflash.inputs.kq_mask_swa ggml_new_tensor_2d(ctx0, mask_type, n_kv_total, GGML_PAD(n_tokens, GGML_KQ_MASK_PAD));即掩码张量形状为{n_kv_total, GGML_PAD(n_tokens, GGML_KQ_MASK_PAD)}——第一个维度覆盖全部 KV 长度第二个维度把 token 数填充到 64 的倍数Vulkan 构建时。类似地DeepSeek 系列图构建中也大量使用该宏如 src/graphs/build_deepseek2.cpp#L682-L710 中关于FA 路径要求 mask 为 F16、连续、且 ne[1] 按 GGML_PAD(n_queries, GGML_KQ_MASK_PAD) 填充的注释以及 src/graphs/build_deepseek4.cpp#L251 中的dsv4_pad_mask_tokens辅助函数。从源码结构可以推断GGML_KQ_MASK_PAD已成为整个图构建层与 Vulkan 后端之间的共享契约图构建负责按对齐粒度分配张量后端负责按对齐粒度无分支消费。六、性能验证RTX 4080 上的 sweep-bench 对比PR 评论区给出了同一张 NVIDIA RTX 4080 上、Q4_0 量化的 Llama-3.1-8B-Instruct、u-batch 1024、Flash Attention 开启时的 sweep-bench 实测数据完整复现如下。Vulkan启用 coopmat2PPTGN_KVT_PP sS_PP t/sT_TG sS_TG t/s102425600.2484128.322.70094.80102425610240.2633887.372.68495.37102425620480.2723769.072.75293.03102425630720.2813639.222.80791.21102425640960.2883560.622.86589.37102425651200.3033380.022.93287.30102425661440.3243158.542.99385.53102425671680.3333074.873.02684.59102425681920.3442977.473.10082.59102425692160.3512920.003.15681.111024256102400.3562876.613.22179.471024256112640.3762725.053.27078.301024256122880.3862651.133.31977.131024256133120.3992564.513.38875.561024256143360.4152470.403.44374.361024256153600.4272400.043.49973.17CUDAPPTGN_KVT_PP sS_PP t/sT_TG sS_TG t/s102425600.1228379.712.054124.65102425610240.1258170.822.092122.39102425620480.1347615.592.154118.84102425630720.1417277.022.221115.26102425640960.1496857.342.290111.77102425651200.1566555.322.371107.97102425661440.1636273.822.412106.14102425671680.1716000.022.467103.77102425681920.1825627.802.527101.32102425692160.1885440.442.58099.231024256102400.1905400.072.66596.041024256112640.2005130.032.70094.831024256122880.2064970.972.75193.061024256133120.2154769.692.81091.101024256143360.2264538.542.86589.341024256153600.2304459.332.93687.18从数据中可以读出几个客观结论PPprompt processing差距最大空 KV 时 CUDA 达 8379.71 t/sVulkan 为 4128.32 t/s差距约 2 倍随着 N_KV 增大两者均衰减但 Vulkan 的相对劣势整体保持例如 N_KV15360 时分别为 4459.33 与 2400.04 t/sTGtoken generation差距较小空 KV 时 CUDA 124.65 t/s vs Vulkan 94.80 t/s约 24% 差距深 KV 时差距进一步缩小87.18 vs 73.17约 16%coopmat2 优于 coopmat1但仍未达到讨论 #562 中20%~25% 差距的预期作者因此将 PP 差距约 2 倍作为后续优化的观察点。需要说明的是这些数据是 PR 提交时作者在特定硬件/驱动/模型组合下的单点实测不代表普遍结论但作为 coopmat2 早期落地时的性能快照具有参考价值。七、如何复现 sweep-bench 测试PR 中使用的工具是仓库自带的 sweep-bench 基准程序源码位于 examples/sweep-bench/sweep-bench.cpp其参数解析复用 common 库的gpt_params。关键参数映射如下-m model.gguf指定模型文件-c N设置 KV 缓存上下文长度sweep 上限-b N设置 batch 大小-ub N设置 ubatch 大小——PP 列即取自params.n_ubatchsweep-bench.cpp#L226-n N设置生成 token 数——TG 列默认取n_predict未指定时为 ubatch/4sweep-bench.cpp#L227-fa开启 Flash AttentionPR 测试即在此模式下进行-ngl N设置 GPU 卸载层数。官方示例命令为./sweep-bench -m model.gguf -c 8192 -b 2048 -ub 512程序会从空 KV 开始按n_ubatch步长逐步填充 KV 缓存并测量每一档的 PP/TG 耗时与吞吐即表中N_KV递增的含义最终以表格或 JSON见 sweep-bench.cpp#L416-L425形式输出。复现 PR 数据时将-ub设为 1024、开启-fa并使用 Q4_0 量化的 Llama-3.1-8B-Instruct 即可对齐表头中的 PP/TG/N_KV 各列。八、总结与延伸阅读PR #574 是一个小而关键的对齐修复它把 Vulkan 后端 Flash Attention 的 KQ mask 填充从 16 改为 64消除了 coopmat2 着色器按 64 行分块读取掩码时的断言冲突并使 ik_llama.cpp 与上游 mainline 的对齐策略保持一致非 Vulkan 后端仍为 16。这一改动也揭示了一个通用工程原则GPU 加速器中张量布局契约shape 对齐、连续性、精度必须与着色器的分块策略严格一致任何一方的演进都可能打破另一方的假设。感兴趣的读者可以继续阅读以下源码与资料对齐宏定义ggml/include/ggml.h#L2486-L2490Vulkan FA 管线构建与断言ggml/src/ggml-vulkan.cpp#L2154-L2173、ggml/src/ggml-vulkan.cpp#L6434-L6435coopmat2 着色器ggml/src/vulkan-shaders/flash_attn_cm2.comp掩码张量分配src/graphs/build_dflash.cpp#L404-L419DeepSeek 系 FA 掩码适配说明src/graphs/build_deepseek2.cpp#L682-L710性能复现工具examples/sweep-bench/sweep-bench.cpp相关讨论Vulkan/ROCm 后端性能与 coopmat 演进见 github-data/discussions/562 - AMD GPU Vulkan _ ROCm_HIP Discussion.md【免费下载链接】ik_llama.cppllama.cpp fork with additional SOTA quants and improved performance项目地址: https://gitcode.com/GitHub_Trending/ik/ik_llama.cpp创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考