커스텀 로짓 프로세서

커스텀 로짓 프로세서 (Custom Logits Processors)

커스텀 로짓 프로세서(custom logits processor)는 vLLM 사용자가 직접 작성해서, vLLM 소스 코드를 수정하거나 재컴파일하지 않고 초기화 시점에 vLLM에 로드되는 프로세서입니다. 내장(built-in) 로짓 프로세서와 반대되는 개념이죠. 이 문서는 커스텀 로짓 프로세서를 작성·로드·사용하는 방법을 보여줍니다.

출처: 문서

중요: 일부 로짓 프로세서 설계는 아직 진행 중이며 API가 가까운 시일 내에 바뀔 수 있습니다. 이 부분의 API 조기 안정화를 목표로 하고 있습니다.

참고: 이 문서는 Model Runner V2(MRV2) 인터페이스를 다룹니다. 레거시 V1 모델 러너(VLLM_USE_V2_MODEL_RUNNER=0)는 vllm.v1.sample.logits_processor.LogitsProcessor라는 다른 인터페이스를 사용하며 여기서는 다루지 않습니다.

본문

로짓 프로세서 배경 (Logits Processors Background)

로짓 프로세서는 다음 토큰 확률 분포를 조정하며, 보통 모델을 원하는 유형의 동작으로 이끄는 것을 목적으로 합니다.

vLLM에서 로짓 프로세서는 **배치 단위(batch granularity)**로 동작합니다. 주어진 엔진 스텝에서 로짓 프로세서는 모델이 출력한 (num_logits_rows) x (vocab_size) 로짓 텐서를 소비합니다. 로짓 행(logits row)은 요청이 아니라는 점에 주의하세요. 행은 매 스텝 재정렬되며, 스펙큘레이티브 디코딩에서는 요청이 드래프트 토큰마다 하나의 행을 소유합니다. 로짓 프로세서는 로짓 텐서의 일부 행에 변환을 적용하고 나머지 행은 수정하지 않습니다. 변환된 로짓 텐서는 샘플링에 사용됩니다.

커스텀 로짓 프로세서 만들기 (Creating a Custom Logits Processor)

커스텀 로짓 프로세서는 vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor를 서브클래싱하고 (최소한) 다음 메서드들을 정의해야 합니다.

• __init__(self, vllm_config: VllmConfig, req_states: LogitsProcRequestState):
  • vllm_config: 엔진 구성 데이터 구조
  • req_states: 지속 배치(persistent batch)의 좁고 읽기 전용 뷰. 디바이스 위 토큰 히스토리(all_token_ids, prompt_len, prefill_len, total_len)와 device, max_num_reqs, vocab_size를 노출합니다.
• add_request(self, req_idx, sampling_params) -> bool:
  • 배치에 들어오는 요청에 대한 슬롯별 상태를 초기화. 슬롯은 free list로 재활용되므로 슬롯별 상태를 여기서 완전히 (재)초기화해야 합니다. freed 슬롯은 절대 읽히지 않으므로 제거 훅(removal hook)은 없습니다.
  • 이 프로세서가 해당 요청의 로짓을 수정하는지 여부를 반환. 샘플러는 어떤 요청도 필요로 하지 않는 배치에 대해 로짓 처리 파이프라인을 건너뜁니다.
• apply_staged_writes(self) -> None (선택):
  • add_request()가 스테이징한 호스트 측 쓰기를 디바이스로 flush. 포워드 패스 전에 스텝당 한 번 호출됩니다.
• validate_params(cls, sampling_params) -> None (선택 classmethod):
  • 잘못된 요청별 인자(특히 커스텀 인자)에 대해 ValueError를 발생. 요청 승인 시점에 실행되므로 유효하지 않은 인자는 샘플러까지 도달하지 않고 요청을 실패시킵니다.
• apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
  • (num_logits_rows) x (vocab_size) 로짓 텐서와 스텝의 배치 레이아웃(ctx: 행-슬롯 매핑, input_ids, 포지션)을 소비.
  • 로짓 프로세서 변환을 배치 단위로 적용.
  • 변환된 로짓 텐서를 반환. 입력 로짓을 제자리(in-place) 또는 비제자리(out-of-place)로 수정할 수 있습니다. in-place가 메모리 효율적입니다.

커스텀 로짓 프로세서에 커스텀 인자 전달하기 (Passing Custom Argument to a Custom Logits Processor)

내장 로짓 프로세서와 달리, 커스텀 로짓 프로세서는 SamplingParams나 vLLM 서버 REST API에 하드코딩되지 않은 구성 인자가 필요할 수 있습니다. 이 문제를 해결하기 위해 커스텀 로짓 프로세서는 vLLM custom arguments 지원을 활용해 사용자로부터 구성 설정을 받을 수 있습니다(SamplingParams의 기존 필드를 사용하는 커스텀 로짓 프로세서를 설계하는 것도 물론 자유롭습니다).

커스텀 로짓 프로세서 예제 구현 (Example Custom Logits Processor Implementation)

아래의 인위적(contrived) 예제는 그것을 활성화한 요청에 대해 target_token 하나를 제외한 모든 토큰을 마스킹하는 커스텀 로짓 프로세서를 구현합니다. target_token을 지정하지 않은 요청에서는 프로세서가 비활성화됩니다. add_request()가 반환값으로 이를 보고하므로, 활성화한 요청이 없을 때 샘플러가 로짓 처리 파이프라인을 완전히 건너뜁니다. 프로세서는 각 요청과 연관된 target_token 커스텀 인자를 SamplingParams.extra_args에서 찾아 활성화 여부와 마스킹하지 않을 토큰을 결정합니다.

예제 커스텀 로짓 프로세서 정의

import torch
from vllm.config import VllmConfig
from vllm.sampling_params import SamplingParams
from vllm.v1.worker.gpu.sample.logits_processor import (
    LogitsProcRequestState,
    LogitsContext,
    LogitsProcessor,
)

class TargetTokenLogitsProcessor(LogitsProcessor):
    """Masks out all tokens except `target_token` (a per-request custom
    argument); requests without it are left alone."""

    def __init__(self, vllm_config: "VllmConfig", req_states: LogitsProcRequestState):
        # Per-slot target; -1 means disabled. Staged on the host and
        # flushed to the device once per step in apply_staged_writes().
        self.target_token = torch.full(
            (req_states.max_num_reqs,), -1, dtype=torch.int64
        )
        self.target_token_dev = torch.full(
            (req_states.max_num_reqs,), -1, dtype=torch.int64, device=req_states.device
        )

    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
        target = (sampling_params.extra_args or {}).get("target_token")
        self.target_token[req_idx] = target if target is not None else -1
        return target is not None

    def apply_staged_writes(self) -> None:
        self.target_token_dev.copy_(self.target_token, non_blocking=True)

    def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
        # ctx.expanded_idx_mapping maps each logits row to its request slot.
        cols = self.target_token_dev[ctx.expanded_idx_mapping.long()]
        rows = torch.nonzero(cols >= 0).squeeze(1)
        if rows.numel() == 0:
            return logits
        kept = logits[rows, cols[rows]].clone()
        logits[rows] = float("-inf")
        logits[rows, cols[rows]] = kept
        return logits

요청별 상태는 요청 슬롯 인덱스로 키가 매겨지고, 슬롯은 free list로 재활용됩니다. 예제는 -1이 비활성화를 뜻하는 슬롯별 상태를 유지합니다. add_request()가 슬롯 항목을 무조건 덮어쓴다는 것에 주목하세요. 이것이 재활용된 슬롯이 요청 간 상태를 누출하지 않게 하는 핵심입니다.

vLLM에 커스텀 로짓 프로세서 로드하기 (Ways to Load Your Custom Logits Processor in vLLM)

로짓 프로세서는 초기화 시점에 로드됩니다. 중요한 점은, vLLM 엔진이 로드를 끝낸 뒤에는 로드된 로짓 프로세서 집합을 수정할 수 없고, 개별 요청에 대해 온디맨드로 새 로짓 프로세서를 로드할 수도 없다는 것입니다.

로드된 클래스는 Model Runner V2 인터페이스에 대해 검증됩니다. vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor를 서브클래싱하지 않은 클래스는 명확한 오류와 함께 거부됩니다.

이 절에서는 로짓 프로세서를 vLLM에 보이게 하고 로드를 트리거하는 여러 방법을 설명합니다.

방법 1: 초기화 시 완전 자격 클래스 이름(FQCN) 전달 (Method 1)

이 방법은 vLLM의 오프라인과 온라인 사용 시나리오 모두에서 지원됩니다. 커스텀 로짓 프로세서의 FQCN(dotted.path.to.module:ClassName 형태)을 LLM과 AsyncLLM Python 생성자에 인자로 넘기거나, vllm serve에 다음 문법으로 CLI 인자로 넘길 수 있습니다.

vllm serve ... --logits_processors <logits processor 1> <logits processor 2> ...

FQCN에 대한 요구 사항은 다음과 같습니다.

• Python의 importlib.import_module()이 FQCN의 dotted 경로 부분을 resolve하여 모듈로 로드할 수 있어야 함.
• FQCN의 클래스 이름 부분이 로드된 모듈에서 import 가능해야 함.
• FQCN이 가리키는 객체가 LogitsProcessor의 서브클래스여야 함.

아래 예시를 보세요.

Python에서 LLM에 FQCN 전달

# Pass in FQCN
llm = LLM(
    model="facebook/opt-125m",
    logits_processors=["your.module.path:TargetTokenLogitsProcessor"],
)

Python에서 AsyncLLM에 FQCN 전달

# Pass in FQCN
engine_args = AsyncEngineArgs(model="facebook/opt-125m",
                              logits_processors=["your.module.path:TargetTokenLogitsProcessor"])
async_llm = AsyncLLM.from_engine_args(engine_args)

CLI로 vLLM 서버에 FQCN 전달

vllm serve facebook/opt-125m --logits_processors your.module.path:TargetTokenLogitsProcessor

방법 2: Python 환경에 엔트리 포인트로 설치된 커스텀 로짓 프로세서 자동 감지 (Method 2)

setuptools는 설치된 패키지가 "엔트리 포인트(entry points)"라는 메타데이터 조각을 통해 다른 Python 프로그램에 플러그인으로 제공되게 할 수 있습니다.

초기화 동안 vLLM은 vllm.logits_processors 엔트리 포인트 그룹을 자동으로 스캔하고 발견한 설치된 로짓 프로세서를 로드합니다.

커스텀 로짓 프로세서를 담은 Python 패키지를 만들었다고 가정해 보세요. 각 로짓 프로세서에 대해 고유한 엔트리포인트를 로짓 프로세서 Python 패키지에 추가하면 vLLM에 노출할 수 있습니다. 아래 예시는 프로젝트의 pyproject.toml에 엔트리포인트를 추가하는 방법입니다.

커스텀 로짓 프로세서를 Python 엔트리포인트로 노출

[project.entry-points."vllm.logits_processors"]
target_token_logits_processor = "your.module.path:TargetTokenLogitsProcessor"

패키지가 설치되면, vLLM이 초기화될 때마다 커스텀 로짓 프로세서가 자동으로 로드됩니다. 로짓 프로세서가 엔트리 포인트로 노출되어 있다면 초기화 시점에 LLM이나 AsyncLLM 생성자 또는 vLLM 서버에 명시적으로 넘길 필요가 없습니다.

참고: vLLM은 vllm.logits_processors 그룹 아래 엔트리포인트로 노출된 모든 로짓 프로세서를 항상 로드합니다.

방법 3 (오프라인 전용): vLLM 생성자에 Python 클래스 객체 전달 (Method 3)

LLM과 AsyncLLM 생성자에 커스텀 로짓 프로세서 클래스 객체 하나 이상을 전달할 수 있습니다. 이 옵션은 매우 유연합니다. 로짓 프로세서 클래스는 (1) LLM이나 AsyncLLM을 인스턴스화하는 같은 Python 소스 파일 안에 로컬로 정의되거나, (2) Python 패키지에서 import될 수 있습니다.

Python에서 LLM 또는 AsyncLLM에 커스텀 로짓 프로세서 클래스 객체 전달

# Import custom logits processor
from some.module import TargetTokenLogitsProcessor

# ...or...

# Define custom logits processor locally
from vllm.v1.worker.gpu.sample.logits_processor import LogitsProcessor

class TargetTokenLogitsProcessor(LogitsProcessor):
    # See TargetTokenLogitsProcessor implementation above
    ...

# Pass class object to LLM constructor
llm = LLM(
    model="facebook/opt-125m",
    logits_processors=[TargetTokenLogitsProcessor],
)

# Pass class object to AsyncLLM constructor
engine_args = AsyncEngineArgs(model="facebook/opt-125m",
                              logits_processors=[TargetTokenLogitsProcessor])
async_llm = AsyncLLM.from_engine_args(engine_args)

요청에 커스텀 로짓 프로세서 호출하기 (Invoking a Custom Logits Processor Against a Request)

커스텀 로짓 프로세서의 설계가, 주어진 요청에 대해 로짓 프로세서를 활성화/비활성화할지, 프로세서 구성에 어떤 인자를 제공할지를 결정합니다.

아래 예시는 TargetTokenLogitsProcessor에 커스텀 인자(target_token)를 전달해 (1) 그 특정 요청에 대해 로짓 프로세서를 활성화하고 (2) 동작을 제어하는 방법을 보여줍니다.

vLLM REST API: 요청에 커스텀 로짓 프로세서 구성

curl http://localhost:8000/v1/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "Qwen/Qwen2.5-1.5B-Instruct",
        ...
        "vllm_xargs": {"target_token": 67}
    }'

OpenAI SDK: 요청에 커스텀 로짓 프로세서 구성

batch = await client.completions.create(
    model="Qwen/Qwen2.5-1.5B-Instruct",
    ...,
    extra_body={
        "vllm_xargs": {
            "target_token": 67
        }
    }
)

오프라인: LLM 요청에 커스텀 로짓 프로세서 구성

outputs_logitproc = llm.generate("your prompt",
                                 SamplingParams(...,
                                    extra_args={"target_token": 67}))

오프라인: AsyncLLM 요청에 커스텀 로짓 프로세서 구성

async for out in engine.generate(request_id="your request id",
                                 prompt="your prompt",
                                 sampling_params=SamplingParams(...,
                                    extra_args={"target_token": 67})):

    # Process async request outputs
    ...

커스텀 로짓 프로세서 작성 모범 사례 (Best Practices for Writing Custom Logits Processors)

• 로짓 프로세서가 배치 단위로 동작한다는 점을 감안해 효율적인 apply()를 작성하세요. 예를 들어 apply()를 효율적인 벡터화 연산으로 구현하고, 요청별 상태를 행별 Python 루프 대신 add_request()에서 연속 텐서로 스테이징할 수 있습니다.
• 무엇을 결정할지는 로짓 프로세서 작성자의 몫입니다.
  • 그 요청에 대해 로짓 프로세서의 동작을 구성하는 요청별 속성. 커스텀 로짓 프로세서의 add_request() 오버라이드가 SamplingParams 필드를 로짓 프로세서 상태로 어떻게 매핑할지를 결정합니다.
  • 요청별로 로짓 프로세서가 활성화되는지 여부의 조건. 커스텀 로짓 프로세서가 항상 모든 요청에 동작하도록 의도한 게 아니라면, 주어진 요청에서 로짓 프로세서를 비활성화할 수 있게 작성하세요. 예를 들어 인자를 기본값 None으로 두거나 0.0 같은 특정 do-nothing 인자 값을 전달하는 방식입니다. 그런 요청에 대해 add_request()에서 False를 반환해, 어떤 요청도 필요로 하지 않는 배치에서 샘플러가 로짓 처리 파이프라인을 건너뛰게 하세요.
• 호스트 측 쓰기를 add_request()에서 스테이징하고 apply_staged_writes()에서 스텝당 한 번 flush하세요. 디바이스 텐서를 직접 쓰지 마세요. 모든 내장 샘플러 상태가 따르는 규율입니다.
• 제거 훅이 없으므로 add_request()는 슬롯 상태를 완전히 (재)초기화해야 합니다. 슬롯의 이전 점유자가 남긴 상태에 의존하지 마세요.

더 알아보기 (Learn more)

• Custom Arguments — 커스텀 로짓 프로세서에 커스텀 인자를 전달하는 vLLM 지원
• SamplingParams API — SamplingParams 필드 참조