vLLM 리랭커로 RAG 검색 재정렬하기

로컬 임베딩으로 RAG 검색을 붙여 돌려보면, 어느 순간 임베딩 모델만으로는 상위 몇 건의 순서가 애매하게 밀리는 걸 본다. 코사인 유사도로 20건을 긁어왔는데 정작 정답 문단이 7등, 8등에 처박혀 있고 top-3만 LLM에 넣으니 답이 틀어지는 식이다. 리랭커를 한 단계 끼우면 이 순서가 꽤 깔끔하게 정리된다. 마침 vLLM이 리랭커를 pooling 러너로 서빙하는 길을 열어둬서, 추론 서버를 하나 더 세우지 않고 같은 스택 안에서 끝낼 수 있다.

결론부터 말하자면, vllm serve <모델> --runner pooling으로 크로스 인코더 리랭커를 띄우고 /rerank 또는 /score에 질의와 후보 문서를 던지면 재정렬된 점수가 돌아온다. bge-reranker-v2-m3는 옵션 없이 바로 뜨고, Qwen3-Reranker는 --hf_overrides로 분류 헤드를 꽂아줘야 한다. 벡터검색이 뽑은 상위 20건을 리랭커로 다시 줄 세운 뒤 top-k를 LLM에 넘기는 구조다.

1. 임베딩이 놓치는 순서를 크로스 인코더가 잡는다

임베딩 기반 벡터검색은 질의와 문서를 각각 벡터로 바꿔 둔 뒤 거리만 잰다. 질의를 받을 때 문서를 다시 읽지 않으니 빠르지만, 질의와 문서가 서로를 보고 상호작용할 기회가 없다. 그래서 단어는 겹치는데 맥락이 어긋난 문단이 상위로 올라오곤 한다.

리랭커는 여기서 쓰는 크로스 인코더다. 질의와 문서를 한 쌍으로 묶어 모델에 같이 넣고, 둘의 관련도를 하나의 점수로 뱉는다. 질의-문서 쌍마다 모델을 한 번씩 돌려야 해서 느리지만, 상호작용을 보기 때문에 순위 정확도가 높다. 그래서 벡터검색으로 후보를 싸게 좁히고, 그 상위 몇십 건만 리랭커로 비싸게 다시 매기는 2단계 구성이 자리 잡았다.

vLLM 쪽 용어로는 리랭커가 pooling 모델에 속한다. 텍스트를 생성하지 않고 입력을 벡터나 점수로 뭉쳐(pool) 내보내는 부류다. 이 중에서도 분류 헤드의 출력 라벨이 1개(num_labels == 1)인 모델만 점수 엔드포인트로 동작한다. 라벨 하나짜리 출력이 곧 관련도 점수가 된다.

2. bge-reranker-v2-m3를 pooling으로 띄우기

가장 손이 덜 가는 건 BAAI의 bge-reranker-v2-m3다. 다국어를 지원하고 한국어 문단도 무난하게 재정렬한다. 별도 override 없이 러너만 pooling으로 지정하면 뜬다.

$ vllm serve BAAI/bge-reranker-v2-m3 --runner pooling --port 8000
INFO ... Using pooling runner for model BAAI/bge-reranker-v2-m3.
INFO ... Supported tasks: ['encode', 'embed', 'classify', 'score']
INFO ... Starting vLLM API server on http://0.0.0.0:8000
INFO ... Route: /score, Methods: POST
INFO ... Route: /v1/score, Methods: POST
INFO ... Route: /rerank, Methods: POST
INFO ... Route: /v1/rerank, Methods: POST
INFO ... Route: /v2/rerank, Methods: POST

--runner의 기본값은 auto라서, 많은 경우 그냥 두면 vLLM이 모델 설정을 보고 pooling으로 잡는다. 다만 아키텍처 자동 판별이 애매한 모델이 있어서 리랭커를 올릴 때는 --runner pooling을 명시하는 편이 덜 헷갈린다.

뜨고 나면 /rerank에 질의 하나와 후보 문서 목록을 넣어 확인한다. 응답은 관련도 순으로 정렬돼 돌아온다.

$ curl -s http://127.0.0.1:8000/v1/rerank \
  -H 'Content-Type: application/json' \
  -d '{
    "model": "BAAI/bge-reranker-v2-m3",
    "query": "vLLM에서 KV 캐시를 CPU로 내리는 옵션",
    "documents": [
      "Route53으로 가중치 기반 라우팅을 설정한다.",
      "swap_space 옵션으로 KV 캐시 일부를 CPU 메모리로 오프로딩한다.",
      "Prometheus에서 GPU 사용률 지표를 수집한다."
    ],
    "top_n": 3
  }' | python3 -m json.tool

{
    "id": "rerank-...",
    "model": "BAAI/bge-reranker-v2-m3",
    "results": [
        {
            "index": 1,
            "document": {"text": "swap_space 옵션으로 KV 캐시 일부를 CPU 메모리로 오프로딩한다."},
            "relevance_score": 0.98
        },
        {
            "index": 0,
            "document": {"text": "Route53으로 가중치 기반 라우팅을 설정한다."},
            "relevance_score": 0.01
        },
        {
            "index": 2,
            "document": {"text": "Prometheus에서 GPU 사용률 지표를 수집한다."},
            "relevance_score": 0.01
        }
    ]
}

점수의 절댓값은 환경과 모델에 따라 다르지만, 정답 문단의 relevance_score가 나머지와 확연히 벌어지는 게 핵심이다. 원래 입력 순서는 index로 추적하니, 벡터검색 결과 배열과 매핑해 두면 재정렬된 순서를 그대로 끌어올 수 있다.

3. Qwen3-Reranker는 override로 분류 헤드를 꽂는다

Qwen3-Reranker는 원래 생성형 구조로 배포돼서, 그대로 올리면 vLLM이 점수 모델로 인식하지 못한다. --hf_overrides로 아키텍처를 Qwen3ForSequenceClassification으로 바꾸고, "no"/"yes" 토큰을 분류 헤드로 끌어다 쓰도록 지정해야 한다. 이 모델은 질의가 문서를 충족하는지 yes/no로 판정하도록 학습됐고, 그 두 토큰의 확률을 점수로 환산하는 방식이다.

$ vllm serve Qwen/Qwen3-Reranker-0.6B \
  --runner pooling \
  --hf_overrides '{"architectures": ["Qwen3ForSequenceClassification"], "classifier_from_token": ["no", "yes"], "is_original_qwen3_reranker": true}' \
  --port 8000

Qwen3-Reranker는 질의와 문서를 감싸는 프롬프트 형식이 정해져 있다. 시스템 지시로 "yes 또는 no로만 답하라"를 주고, <Instruct>, <Query>, <Document> 틀에 맞춰 넣는다. 이 포장을 빼먹으면 점수가 뭉개진다. 공식 예제는 prefix/suffix를 직접 붙여 요청을 만든다.

prefix = ('<|im_start|>system\n'
          'Judge whether the Document meets the requirements based on the Query '
          'and the Instruct provided. Note that the answer can only be "yes" or "no".'
          '<|im_end|>\n<|im_start|>user\n')
suffix = '<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n'

q_tmpl = '{prefix}<Instruct>: {inst}\n<Query>: {query}\n'
d_tmpl = '<Document>: {doc}{suffix}'

resp = requests.post('http://127.0.0.1:8000/v1/rerank', json={
    'query': q_tmpl.format(prefix=prefix, inst='Given a web search query, retrieve relevant passages that answer the query', query='vLLM에서 KV 캐시를 CPU로 내리는 옵션'),
    'documents': [d_tmpl.format(doc=d, suffix=suffix) for d in docs],
}).json()

0.6B 외에 4B, 8B 체급이 있다. 0.6B는 GPU 메모리 몇 GB로 돌아가서 임베딩 서버 옆에 같이 띄우기 좋고, 한국어 품질이 더 필요하면 4B로 올려본다. 체급을 키울수록 쌍당 추론 비용이 늘어나니 후보 수(top-n)를 과하게 잡지 않는 게 좋다.

4. 벡터검색 상위 후보를 리랭커에 흘려보내기

실제 RAG 코드에선 벡터검색이 먼저 상위 K건을 뽑고, 그 텍스트만 리랭커로 넘긴다. 아래는 pgvector나 임베딩 API로 받은 후보 리스트를 /v1/rerank에 넣어 상위 몇 건만 남기는 형태다.

import requests

# 1단계: 벡터검색으로 후보 20건 (candidates = [{'id':.., 'text':..}, ...])
docs = [c['text'] for c in candidates]

# 2단계: 리랭커로 재정렬
r = requests.post('http://127.0.0.1:8000/v1/rerank', json={
    'model': 'BAAI/bge-reranker-v2-m3',
    'query': user_query,
    'documents': docs,
    'top_n': 5,
}).json()

# 3단계: 상위 5건만 LLM 컨텍스트로
top = [candidates[item['index']] for item in r['results']]
for item in r['results'][:5]:
    print(round(item['relevance_score'], 3), candidates[item['index']]['id'])

후보 수는 20에서 50 사이로 두고 top_n은 3에서 5로 좁히는 구성이 흔하다. 벡터검색이 정답을 상위 50 안에만 넣어주면 리랭커가 그 안에서 위로 끌어올리는 역할을 한다. 반대로 벡터검색이 50건 안에 정답을 아예 못 넣으면 리랭커도 못 살린다. 그래서 1단계 recall을 먼저 확인하고 2단계를 붙여야 한다.

점수 쌍 하나하나가 필요하면 /score를 쓴다. /rerank가 정렬된 결과를 주는 반면, /score는 질의-문서 쌍의 점수를 입력 순서대로 data 배열에 돌려준다. 평가 스크립트에서 모든 쌍의 점수를 뽑아 지표를 계산할 때 편하다.

$ curl -s http://127.0.0.1:8000/score \
  -H 'Content-Type: application/json' \
  -d '{"model": "BAAI/bge-reranker-v2-m3",
       "text_1": "vLLM에서 KV 캐시를 CPU로 내리는 옵션",
       "text_2": ["swap_space로 KV 캐시를 CPU에 오프로딩한다.", "Route53 가중치 라우팅"]}' \
  | python3 -m json.tool
{
    "data": [
        {"index": 0, "score": 0.97},
        {"index": 1, "score": 0.01}
    ]
}

5. 적중률은 같은 질의셋으로 전후를 비교한다

리랭커가 값을 하는지는 숫자로 봐야 한다. 정답 문서가 라벨링된 질의셋을 하나 만들어 두고, 벡터검색만 썼을 때와 리랭커를 붙였을 때의 상위 k 적중률을 같은 셋으로 재면 된다. Hit@k는 상위 k건 안에 정답이 들었는지, MRR은 정답이 몇 등에 왔는지를 본다.

# eval_set = [{'query':.., 'gold_id':.., 'candidates':[...]}, ...]
def hit_at_k(order, gold_id, k):
    return gold_id in [c['id'] for c in order[:k]]

base_hit = rerank_hit = 0
for row in eval_set:
    base = row['candidates']                 # 벡터검색 순서 그대로
    r = rerank(row['query'], base)           # 리랭커 통과
    base_hit   += hit_at_k(base, row['gold_id'], 3)
    rerank_hit += hit_at_k(r,    row['gold_id'], 3)

n = len(eval_set)
print('Hit@3  base  :', base_hit / n)
print('Hit@3  rerank:', rerank_hit / n)

개선 폭은 데이터와 모델에 따라 크게 갈려서 일반화하기 어렵다. 임베딩 모델이 약하거나 문서가 서로 비슷해 상위권이 뭉개지는 셋일수록 리랭커의 Hit@3 상승이 두드러지고, 임베딩만으로 이미 상위가 잘 정렬된 셋에서는 차이가 작을 수 있다. 중요한 건 추측이 아니라 자기 질의셋으로 이 두 줄 숫자를 직접 비교하는 것이다. Ragas 같은 도구로 컨텍스트 적합도까지 같이 보면 더 깔끔하다.

6. 올리다 걸리는 지점들

Qwen3-Reranker를 --hf_overrides 없이 올리면 생성형 러너로 잡혀 /score, /rerank 라우트가 아예 안 생긴다. 서버 로그의 Supported tasks에 score가 없고 generate만 보이면 override가 안 먹은 것이다. JSON 따옴표가 셸에서 깨지는 경우가 많으니 작은따옴표로 전체를 감싸고 안쪽은 큰따옴표를 쓴다.

긴 문서를 넣으면 모델 최대 길이를 넘겨 뒷부분이 잘린다. --max-model-len을 질의와 문서 합산 길이에 맞게 올리거나, 요청에 truncate_prompt_tokens를 줘서 자르는 지점을 통제한다. 후보 문서가 과하게 길면 재정렬 자체가 왜곡되니, 청크 단위를 리랭커 입력 길이에 맞춰 잘라 두는 게 낫다.

후보를 너무 많이 넘기면 쌍당 추론이 선형으로 늘어 지연이 커진다. GPU 한 장에서 임베딩 서버와 리랭커를 같이 돌리는 구성이라면 top_n과 후보 수를 먼저 보수적으로 잡고, 적중률이 떨어지지 않는 선에서 후보 수를 줄여 지연을 맞춘다. 임베딩과 리랭커를 각각 포트를 나눠 띄우면 지표 수집과 스케일링이 따로 돌아 운영이 편하다.