llama.cpp reranking源码分析

大模型时代,reranker一直是提高RAG有效性的重要工具。相对于初筛阶段向量检索,精排阶段的reranker需要query和每个候选document做相关计算。初筛已经将候选documents限制在一个相对较小范围,但依然要进行大量的相关性计算。

llama.cpp是广泛使用的模型量化工具,支持16、8、4、2位模型量化,降低存储占用,提高运行效率。而且llama.cpp的server模式,兼容openai格式的访问接口。

这里通过阅读llama.cpp reranking源码,分析llama.cpp运行reranker方式,探索可能的优化点。

1 llama.cpp reranking

如下所示,目前(2025.8.26)llama.cpp依然采用较传统的reranking计算方法。

llama.cpp收到query和documents后,先tokenize处理query,后处理documents。

然后,在循环体中,针对每个处理后的tokenized_docs[i],和query一起送入tasks队列进行计算。

query tokenize处理

llama_tokens tokenized_query = tokenize_input_prompts(ctx_server.vocab, query, /* add_special */ false, true)[0];

{

documents tokenize处理

auto tokenized_docs = tokenize_input_prompts(ctx_server.vocab, documents, /* add_special */ false, true);

for (size_t i = 0; i < tokenized_docs.size(); i++) {

针对每个 tokenized_docs[i]进行reranking计算

auto tmp = format_rerank(ctx_server.vocab, tokenized_query, tokenized_docs[i]);

...

}

可见,llama.cpp采用串行方式分别计算每个<query, document>对的相关性。

以下是llama.cpp处理reranking的完整代码,来源如下

https://github.com/ggml-org/llama.cpp/blob/master/tools/server/server.cpp

复制代码
        llama_tokens tokenized_query = tokenize_input_prompts(ctx_server.vocab, query, /* add_special */ false, true)[0];

        // create and queue the task
        json responses = json::array();
        bool error = false;
        std::unordered_set<int> task_ids;
        {
            std::vector<server_task> tasks;
            auto tokenized_docs = tokenize_input_prompts(ctx_server.vocab, documents, /* add_special */ false, true);
            tasks.reserve(tokenized_docs.size());
            for (size_t i = 0; i < tokenized_docs.size(); i++) {
                auto tmp = format_rerank(ctx_server.vocab, tokenized_query, tokenized_docs[i]);
                server_task task   = server_task(SERVER_TASK_TYPE_RERANK);
                task.id            = ctx_server.queue_tasks.get_new_id();
                task.index         = i;
                task.prompt_tokens = server_tokens(tmp, ctx_server.mctx != nullptr);
                tasks.push_back(std::move(task));
            }

            task_ids = server_task::get_list_id(tasks);
            ctx_server.queue_results.add_waiting_tasks(tasks);
            ctx_server.queue_tasks.post(std::move(tasks));
        }

        ctx_server.receive_multi_results(task_ids, [&](std::vector<server_task_result_ptr> & results) {
            for (auto & res : results) {
                GGML_ASSERT(dynamic_cast<server_task_result_rerank*>(res.get()) != nullptr);
                responses.push_back(res->to_json());
            }
        }, [&](const json & error_data) {
            res_error(res, error_data);
            error = true;
        }, req.is_connection_closed);

        if (error) {
            return;
        }

        // write JSON response
        json root = format_response_rerank(
            body,
            responses,
            is_tei_format,
            documents);

        res_ok(res, root);

2 reranking优化探索

如上所述,llama.cpp采用串行方式分别计算每个<query, document>对的相关性。

其实,目前已经有一些效率更高的reranker计算方式。

比如,采用batch,一次性并行计算多个<query, document>对。

另外,针对reranker计算中query和instruct重复计算问题,采用自定义注意力掩码方式,在推理中共享query和instruct部分,仅计算documents部分。

实现可参考基于自定义注意力掩码的reranker运行速度优化-CSDN博客

reference


llama.cpp http server

https://github.com/ggml-org/llama.cpp/blob/master/tools/server/server.cpp

基于llama.cpp的量化版reranker模型调用示例

https://blog.csdn.net/liliang199/article/details/150619898

基于自定义注意力掩码的reranker运行速度优化

https://blog.csdn.net/liliang199/article/details/150612837

相关推荐
喜欢吃豆19 小时前
掌握本地化大语言模型部署:llama.cpp 工作流与 GGUF 转换内核全面技术指南
人工智能·语言模型·架构·大模型·llama·llama.cpp·gguf
illuspas3 天前
Ubuntu 24.04下编译支持ROCm加速的llama.cpp
linux·ubuntu·llama
缘友一世5 天前
LLama3架构原理浅浅学学
人工智能·自然语言处理·nlp·transformer·llama
我们没有完整的家6 天前
批量吞吐量实测:Llama-2-7b 昇腾 NPU 六大场景数据报告
llama
asfdsfgas7 天前
从加载到推理:Llama-2-7b 昇腾 NPU 全流程性能基准
人工智能·llama
asdfsdgss7 天前
FP16 vs INT8:Llama-2-7b 昇腾 NPU 精度性能基准报告
llama
猿代码_xiao7 天前
大模型微调完整步骤( LLama-Factory)
人工智能·深度学习·自然语言处理·chatgpt·llama·集成学习
wei_shuo9 天前
Llama-2-7b 昇腾 NPU 测评总结:核心性能数据、场景适配建议与硬件选型参考
大模型·llama·昇腾
凯子坚持 c9 天前
Llama-2-7b在昇腾NPU上的六大核心场景性能基准报告
java·开发语言·llama
落798.9 天前
【在昇腾NPU上部署Llama-2-7B:从环境配置到性能测试的完整实战】
经验分享·llama·1024程序员节