llama.cpp b11372重写lightning indexer评分路径,索引器分数内存减半
该版本按PR #29825 的提交说明,将qwen4exp索引器逐头计算并就地累加,使同时存活的f32分数张量由两个降为一个;CUDA、Metal、Vulkan后端同步调整了对lightning indexer的头数支持。
AI解读:这次改动针对的是qwen4exp索引器在长上下文下的内存占用:提交说明称原实现一次性对所有头做乘积并整流一份副本,导致两个 [n_pool, n_idx_h, n_tokens] 的f32张量同时存活,是图中最大的缓冲区。
新实现改为每个头单独做乘积,整流后就地累加到单个 [n_pool, n_tokens] 分数上,因此同时活跃的大张量从一个减为另一个的规模,即标题所说的“减半”。这一说法来自PR提交说明,不是第三方实测。
评审意见要求改用普通ggml_add与ggml_relu;提交说明称图分配器在源张量没有其他消费者时已会就地执行,因此不需要 _inplace变体,计算缓冲和速度不变。
索引器后来改为调用ggml_lightning_indexer:提交说明解释,整流后头分数的无权重和乘以 1/sqrt(head_dim),等价于所有头权重都设为该缩放值的lightning indexer,键只被读一次,不再实体化逐头分数。
后端层面,CUDA将 4 头派发到向量内核(对wmma tile来说头数太少)并在supports_op中接受,test-backend-ops覆盖 4 头;Metal将lightning indexer头数作为函数常量读取并对最后一个头tile补零,声明任意头数可运行且 64 头不变。
Vulkan改为按 64 个键对 8 个token分块,键在共享内存中暂存一次、查询逐头加载,每个调用负责一个键的两个token,无需跨调用归约;子组变体和扁平派发被移除,网格为keys x tokens x streams,并向量化加载、使用fp16点积。这些均是提交说明中的实现描述,未给出长上下文下的吞吐或显存实测数字。
llama.cpp发布b11372版本,其中包含PR #29825,标题为“qwen4exp : halve the indexer score memory”。按提交说明,改动使qwen4exp索引器在长上下文下同时存活的评分张量由两份减为一份。
索引器评分内存如何减半
提交说明称,原实现一次性对所有头做乘积,并整流其副本,因此有两个 [n_pool, n_idx_h, n_tokens] 的f32张量同时存活,在长上下文时是图中最大的缓冲区。
新实现让每个头各自做乘积,整流后就地累加进一个 [n_pool, n_tokens] 的分数。提交说明将这项收益表述为内存减半,未附带显存或吞吐的实测数字。
评审后改用普通算子与lightning indexer
提交说明记录了一轮评审修改:按CISC意见,在索引器头循环中使用普通的ggml_add和ggml_relu。理由是图分配器在源张量没有其他消费者时已经会就地执行,不需要 _inplace变体,计算缓冲和速度不变。
另一轮按am17an意见的改动称,整流后头分数的无权重和乘以 1/sqrt(head_dim),等价于把所有头权重都设为该缩放值的lightning indexer,于是索引器在池化键上以f16 pool mask调用ggml_lightning_indexer;键对所有头只读一次,不再实体化逐头分数。
CUDA、Metal、Vulkan对头数的处理
CUDA方面,提交说明称 4 头被派发到向量内核,因为对wmma tile来说头数太少;supports_op接受该情形,test-backend-ops覆盖 4 头。
Metal方面,内核从函数常量读取头数并对最后一个头tile补零,提交说明称任意头数都能运行,64 头行为不变。
Vulkan方面,一个工作组对 8 个token计算 64 个键的分数:键在共享内存中暂存一次,查询逐头加载,每个调用拥有一个键对应两个token,因此没有点积需要跨调用归约。子组变体和扁平派发被移除,网格为keys x tokens x streams;提交说明还提到向量化加载并使用fp16点积。