大模型時代,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 taskjson 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 responsejson 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