커스텀 인자 파싱

커스텀 인자 파싱 (Custom Arguments)

vLLM은 사용자 정의 기능 모음(예: 커스텀 로짓 프로세서, 커스텀 출력 프로세서 또는 사전 처리)을 특별한 console script 엔트리포인트로 등록할 수 있어요. vLLM은 이렇게 등록된 커스텀 코드에 대해 사용자 정의 인자(custom arguments)를 파싱해요. 사용자는 --custom.boolean_flag, --custom.key=value, --custom.json_option='{"foo": "bar"}' 같은 문법으로 접근할 수 있어요. 이 기능의 구현 세부 사항은 아래에서 자세히 설명할게요.

console script 엔트리포인트

vLLM에게 알려주고 싶은 코드를 등록하려면 pyproject.toml의 [project.scripts] 테이블에 엔트리포인트를 추가하면 돼요. 예를 들어 새 커스텀 로짓 프로세서를 등록한다면, 먼저 파이썬 클래스를 정의하고 그것을 향하는 엔트리포인트를 추가해요.

파이썬 클래스 정의

선택 사항으로 커스텀 로짓 프로세서 클래스 내에서 글로벌 타입 힌트 변수와 데코레이터를 정의할 수 있어요.

vllm.custom_arg_parser.DECLARE_FLAG: vLLM 커맨드라인 프로그램이 이 타입을 사용하는 값을 위한 부울 플래그를 만들게 하는 데코레이터 (기본값 True로 시작)
vllm.custom_arg_parser.DECLARE_KEY_VALUE: vLLM 커맨드라인이 키-값 쌍 컬렉션 인자(--custom.key=value)를 만들게 하는 데코레이터. 값은 :로 구분된 다른 키-값 쌍이거나 파이썬 리스트일 수 있어요
vllm.custom_arg_parser.DECLARE_JSON: vLLM 커맨드라인 프로그램이 JSON 문자열(--custom.json_option='{"foo": "bar"}')을 파싱하게 하는 데코레이터

vllm.custom_arg_parser.OPTIONS에 글로벌 변수를 선언하면 위에서 정의한 타입 힌트들을 커스텀 코드를 위한 옵션으로 모두 수집해요. 이것은 vLLM에게 커스텀 코드가 필요로 하는 사용자 지정 인자를 알려주는 핵심이에요.

# custom_code/custom_logits_processor.py

import torch
from typing import Literal, TypedDict, Optional

from vllm import SamplingParams
from vllm.custom_arg_parser import DECLARE_FLAG, DECLARE_KEY_VALUE, DECLARE_JSON, OPTIONS
from vllm.v1.sample.logits_processor import LogitsProcessor

@DECLARE_FLAG
NO_SPECIAL_TOKENS: bool = True  # vLLM 프로그램이 --custom.no-special-tokens 부울 플래그를 만들어요

@DECLARE_KEY_VALUE
TOKEN_BIASES: dict[str, float] = {}

@DECLARE_JSON
TOKEN_CONSTRAINTS: dict[str, list[str]] = {}

class CustomLogitsProcessor(LogitsProcessor):
    """vLLM에 커스텀 로짓 프로세서를 등록하는 예제"""
    def __init__(self, vllm_config, device, is_pin_memory) -> None:
        super().__init__()
        # ...

    def apply(self, logits: torch.Tensor) -> torch.Tensor:
        # ...
        return logits

    def is_argmax_invariant(self) -> bool:
        return True

    def update_state(self, batch_update) -> None:
        return None

    @classmethod
    def validate_params(cls, sampling_params: SamplingParams) -> None:
        # ...
        return None

    @classmethod
    def parse_custom_arguments(cls):
        return {
            "no_special_tokens": NO_SPECIAL_TOKENS,
            "token_biases": TOKEN_BIASES,
            "token_constraints": TOKEN_CONSTRAINTS,
        }

예제에서는 parse_custom_arguments 클래스 메서드를 정의해서, 등록된 타입 힌트들이 파싱된 후 커스텀 코드에서 접근할 수 있게 해요.

pyproject.toml 등록 예제

[project.scripts]
vllm-logits-processor = "custom_code.custom_logits_processor:CustomLogitsProcessor"

이 엔트리포인트는 vLLM이 vLLM CLI 프로그램을 시작할 때 커스텀 코드를 자동으로 로드하게 해요. 모든 vLLM console scripts는 동일한 커스텀 로짓 프로세서를 인식하도록 조정돼 있어요.

커스텀 로짓 프로세서가 로드된 것을 확인할 수 있는 로그는 문서의 해당 섹션에서 확인할 수 있어요.

커스텀 인자 사용법

플래그 인자 (Boolean Flags)

python -m vllm.entrypoints.openai.api_server \
    --custom.no-special-tokens

이 인자는 클래스 정의에서 NO_SPECIAL_TOKENS: bool = True를 사용해 만들어진 --custom.no-special-tokens 플래그를 켭니다.
--custom.no-special-tokens=False 형태로 명시적으로 끌 수도 있어요.

키-값 인자 (Key-Value)

python -m vllm.entrypoints.openai.api_server \
    --custom.token-biases='<|im_start|>:-2.0,<|im_end|>:-5.0'

TOKEN_BIASES: dict[str, float] = {}를 위해 --custom.token-biases 키-값 옵션을 만들어요.
콜론(:)으로 구분된 키-값 쌍을 쉼표로 구분해 여러 개 전달할 수 있어요.
하나 이상의 값을 + 로 구분해 .(점) 초과... 값이 리스트가 되도록 할 수도 있어요.
문서에 예시가 있지만, 실제로는 리스트 형태 파싱도 지원해요.

JSON 인자 (JSON)

python -m vllm.entrypoints.openai.api_server \
    --custom.token-constraints='{"<|im_start|>": ["-2.0", "-5.0"]}'

TOKEN_CONSTRAINTS: dict[str, list[str]] = {}를 위한 --custom.token-constraints JSON 옵션을 만들어요.
값은 유효한 JSON 문자열이어야 해요.

이 커스텀 인자들을 CLI에서 사용하는 전체 예시는 문서의 서빙 예제 섹션에서 확인할 수 있어요.

일반적인 안티패턴과 함정

커스텀 인자 파싱 시스템과 관련된 몇 가지 주의 사항이 있어요:

기본값에 복잡한 타입을 쓰는 경우: @DECLARE_* 데코레이터는 기본 설정으로 정확한 타입이 필요해요. 타입 힌트가 정확하지 않으면 vLLM이 그에 맞는 파서를 만들 수 없어요.
--custom.greeting=world 이름 충돌: 인자 이름 충돌이 있으면 vLLM이 오류를 발생시켜요.
리스트/고급 구문은 문서의 "사용자 정의 인자 구문 (Custom argument values)" 섹션을 참고하세요.

전반적으로, 사용자가 정의한 커스텀 인자를 파싱하는 시스템이라서, 문서화된 사용법을 따르면 문제가 없어요.

더 알아보기

vLLM 공식 문서: Custom Arguments