ik_llama.cpp PR 574 技术解析：将 KQ mask 填充对齐从 16 提升到 64，以适配 Vulkan coopmat2 闪存注意力

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 2（coopmat2）时触发的断言失败，并同步给出了同一张 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 processing（PP）仍比 CUDA 慢约 2 倍。

二、背景知识：什么是 KQ mask，为什么要填充

在自回归 Transformer 的解码与训练中，注意力分数矩阵需要叠加掩码（mask），用来屏蔽未来的 token（causal 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，即真实 batch（token 数）经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 路径特有的硬性需求，而非所有后端通用的改动。

三、为什么是 64：coopmat2 着色器的分块约束

要理解"为什么必须是 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 中可以看到它们的创建逻辑：

1. FA_SCALAR：标量路径，无条件创建，支持 F16、Q4_0、Q8_0；
2. FA_COOPMAT1：基于VK_KHR_cooperative_matrix扩展的 cooperative matrix 1，需设备报告coopmat1_fa_support；
3. FA_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（启用 coopmat2）

PP	TG	N_KV	T_PP s	S_PP t/s	T_TG s	S_TG t/s
1024	256	0	0.248	4128.32	2.700	94.80
1024	256	1024	0.263	3887.37	2.684	95.37
1024	256	2048	0.272	3769.07	2.752	93.03
1024	256	3072	0.281	3639.22	2.807	91.21
1024	256	4096	0.288	3560.62	2.865	89.37
1024	256	5120	0.303	3380.02	2.932	87.30
1024	256	6144	0.324	3158.54	2.993	85.53
1024	256	7168	0.333	3074.87	3.026	84.59
1024	256	8192	0.344	2977.47	3.100	82.59
1024	256	9216	0.351	2920.00	3.156	81.11
1024	256	10240	0.356	2876.61	3.221	79.47
1024	256	11264	0.376	2725.05	3.270	78.30
1024	256	12288	0.386	2651.13	3.319	77.13
1024	256	13312	0.399	2564.51	3.388	75.56
1024	256	14336	0.415	2470.40	3.443	74.36
1024	256	15360	0.427	2400.04	3.499	73.17

CUDA

PP	TG	N_KV	T_PP s	S_PP t/s	T_TG s	S_TG t/s
1024	256	0	0.122	8379.71	2.054	124.65
1024	256	1024	0.125	8170.82	2.092	122.39
1024	256	2048	0.134	7615.59	2.154	118.84
1024	256	3072	0.141	7277.02	2.221	115.26
1024	256	4096	0.149	6857.34	2.290	111.77
1024	256	5120	0.156	6555.32	2.371	107.97
1024	256	6144	0.163	6273.82	2.412	106.14
1024	256	7168	0.171	6000.02	2.467	103.77
1024	256	8192	0.182	5627.80	2.527	101.32
1024	256	9216	0.188	5440.44	2.580	99.23
1024	256	10240	0.190	5400.07	2.665	96.04
1024	256	11264	0.200	5130.03	2.700	94.83
1024	256	12288	0.206	4970.97	2.751	93.06
1024	256	13312	0.215	4769.69	2.810	91.10
1024	256	14336	0.226	4538.54	2.865	89.34
1024	256	15360	0.230	4459.33	2.936	87.18

从数据中可以读出几个客观结论：

1. PP（prompt processing）差距最大：空 KV 时 CUDA 达 8379.71 t/s，Vulkan 为 4128.32 t/s，差距约 2 倍；随着 N_KV 增大，两者均衰减，但 Vulkan 的相对劣势整体保持（例如 N_KV=15360 时分别为 4459.33 与 2400.04 t/s）；
2. TG（token generation）差距较小：空 KV 时 CUDA 124.65 t/s vs Vulkan 94.80 t/s（约 24% 差距），深 KV 时差距进一步缩小（87.18 vs 73.17，约 16%）；
3. 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_ubatch（sweep-bench.cpp#L226）；
• -n N：设置生成 token 数——TG 列默认取n_predict，未指定时为 ubatch/4（sweep-bench.cpp#L227）；
• -fa：开启 Flash Attention（PR 测试即在此模式下进行）；
• -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-L2490
• Vulkan FA 管线构建与断言：ggml/src/ggml-vulkan.cpp#L2154-L2173、ggml/src/ggml-vulkan.cpp#L6434-L6435
• coopmat2 着色器：ggml/src/vulkan-shaders/flash_attn_cm2.comp
• 掩码张量分配：src/graphs/build_dflash.cpp#L404-L419
• DeepSeek 系 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