플러그인 시스템

플러그인 시스템 (Plugin System)

커뮤니티에서 vLLM에 커스텀 기능을 붙이고 싶다는 요청이 아주 많아요. vLLM에는 이런 요구를 담아, vLLM 코드베이스를 수정하지 않고도 사용자가 커스텀 기능을 추가할 수 있게 해주는 플러그인 시스템이 포함돼 있어요. 이 문서에서는 플러그인이 vLLM에서 어떻게 동작하는지, 그리고 플러그인을 어떻게 만들어야 하는지를 설명할게요.

vLLM에서 플러그인이 동작하는 방식

플러그인은 사용자가 등록한 코드로, vLLM이 그걸 실행해요. vLLM의 아키텍처(아키텍처 개요 참고)상 다양한 병렬 기법을 쓰는 분산 추론에서는 여러 프로세스가 관여할 수 있어요. 플러그인을 제대로 동작시키려면 vLLM이 만드는 모든 프로세스가 플러그인을 로드해야 해요. 이 작업을 vllm.plugins 모듈의 load_plugins_by_group 함수가 담당해요.

vLLM이 플러그인을 찾는 방식

vLLM의 플러그인 시스템은 표준 파이썬 entry_points 메커니즘을 사용해요. 이 메커니즘은 개발자가 자기 파이썬 패키지에 함수를 등록해서 다른 패키지가 사용할 수 있게 해줘요. 플러그인 예시:

# inside `setup.py` file
from setuptools import setup

setup(name='vllm_add_dummy_model',
    version='0.1',
    packages=['vllm_add_dummy_model'],
    entry_points={
        'vllm.general_plugins':
        ["register_dummy_model = vllm_add_dummy_model:register"]
    })

# inside `vllm_add_dummy_model/__init__.py` file
def register():
    from vllm import ModelRegistry

    if "MyLlava" not in ModelRegistry.get_supported_archs():
        ModelRegistry.register_model(
            "MyLlava",
            "vllm_add_dummy_model.my_llava:MyLlava",
        )

패키지에 entry point를 추가하는 자세한 방법은 공식 문서를 참고하면 돼요.

모든 플러그인은 세 부분으로 이뤄져요.

  1. 플러그인 그룹 (Plugin group): entry point 그룹의 이름이에요. vLLM은 vllm.general_plugins 그룹으로 일반 플러그인을 등록해요. setup.py에서 entry_points의 key가 바로 이 값이에요. vLLM 일반 플러그인에는 항상 vllm.general_plugins를 써야 해요.
  2. 플러그인 이름 (Plugin name): entry_points 딕셔너리 안의 value(이름)예요. 위 예제에서 플러그인 이름은 register_dummy_model이에요. VLLM_PLUGINS 환경 변수로 이름별로 플러그인을 필터링할 수 있어요. 특정 플러그인만 로드하려면 VLLM_PLUGINS에 그 이름을 설정하면 돼요.
  3. 플러그인 값 (Plugin value): 플러그인 시스템에 등록할 함수나 모듈의 정규화된 이름(FQN)이에요. 위 예제에서 플러그인 값은 vllm_add_dummy_model:register로, vllm_add_dummy_model 모듈 안의 register 함수를 가리켜요.

지원되는 플러그인 유형

  • 일반 플러그인 (그룹 vllm.general_plugins): 주 용도는 커스텀·트리 밖(out-of-the-tree) 모델을 vLLM에 등록하는 거예요. 플러그인 함수 안에서 ModelRegistry.register_model을 호출해서 모델을 등록해요. 공식 모델 플러그인 예시로는 BartForConditionalGeneration 지원을 추가하는 bart-plugin이 있어요.

  • 플랫폼 플러그인 (그룹 vllm.platform_plugins): 커스텀·트리 밖 플랫폼을 vLLM에 등록하는 용도예요. 플러그인 함수는 현재 환경에서 플랫폼을 지원하지 않으면 None을, 지원하면 플랫폼 클래스의 정규화된 이름을 반환해요.

  • IO 프로세서 플러그인 (그룹 vllm.io_processor_plugins): pooling 모델의 프롬프트 입력과 출력 전·후처리를 커스텀으로 등록하는 용도예요. 플러그인 함수는 IOProcessor 클래스의 정규화된 이름을 반환해요.

  • 통계 로거 플러그인 (그룹 vllm.stat_logger_plugins): 커스텀·트리 밖 로거를 vLLM에 등록하는 용도예요. entry point는 StatLoggerBase를 상속한 클래스여야 해요.

  • 엔드포인트 플러그인 (그룹 vllm.endpoint_plugins): OpenAI 호환 API 서버에 커스텀·트리 밖 HTTP 라우트를 등록하는 용도예요. 다른 플러그인 그룹과 달리 엔드포인트 플러그인은 API 서버 프런트엔드 프로세스에서만 로드되고 기본으로는 로드되지 않아요. 인터페이스는 Endpoint Plugins를, opt-in·신뢰 모델은 보안을 참고하세요.

플러그인 작성 가이드라인

  • 재진입성 (re-entrant): entry point에 지정한 함수는 재진입 가능해야 해요. 즉 여러 번 호출돼도 문제가 없어야 해요. 일부 프로세스에서 함수가 여러 번 호출될 수 있기 때문이에요.

플랫폼 플러그인 가이드라인

  1. 플랫폼 플러그인 프로젝트를 만들어요(예: vllm_add_dummy_platform). 프로젝트 구조는 이렇게 생겨야 해요.

    vllm_add_dummy_platform/
├── vllm_add_dummy_platform/
│   ├── __init__.py
│   ├── my_dummy_platform.py
│   ├── my_dummy_worker.py
│   ├── my_dummy_attention.py
│   ├── my_dummy_device_communicator.py
│   ├── my_dummy_custom_ops.py
├── setup.py

  2. setup.py 파일에 다음 entry point를 추가해요.

    setup(
    name="vllm_add_dummy_platform",
    ...
    entry_points={
        "vllm.platform_plugins": [
            "my_dummy_platform = vllm_add_dummy_platform:register"
        ]
    },
    ...
)

    vllm_add_dummy_platform:register가 호출 가능한 함수이고 플랫폼 클래스의 정규화된 이름을 반환하는지 확인하세요. 예:

    def register():
    return "vllm_add_dummy_platform.my_dummy_platform.MyDummyPlatform"

  3. my_dummy_platform.py에 플랫폼 클래스 MyDummyPlatform을 구현해요. 이 클래스는 vllm.platforms.interface.Platform을 상속해야 해요. 인터페이스를 따라 함수를 하나씩 구현하면 돼요. 최소한 구현해야 할 중요한 함수·속성은:

    • _enum: PlatformEnum의 장치 열거형이에요. 보통 PlatformEnum.OOT(out-of-tree)를 써요.
    • device_type: pytorch가 쓰는 장치 타입을 반환해야 해요. 예: "cpu", "cuda" 등
    • device_name: 보통 device_type과 같게 설정해요. 주로 로깅용이에요.
    • check_and_update_config: vLLM 초기화 과정에서 아주 일찍 호출돼요. 플러그인이 vllm 구성을 갱신하는 데 쓰여요. 예를 들어 블록 크기, 그래프 모드 설정 등을 이 함수에서 갱신할 수 있어요. 가장 중요한 건 worker_cls를 이 함수에서 설정해서 vLLM이 워커 프로세스에 어떤 워커 클래스를 쓸지 알게 하는 거예요.
    • get_attn_backend_cls: 어텐션 백엔드 클래스의 정규화된 이름을 반환해야 해요.
    • get_device_communicator_cls: 장치 커뮤니케이터 클래스의 정규화된 이름을 반환해야 해요.

  4. my_dummy_worker.py에 워커 클래스 MyDummyWorker을 구현해요. WorkerBase를 상속해야 해요. 기본 클래스의 인터페이스는 vLLM 곳곳에서 호출되므로 기본적으로 전부 구현해야 해요. 모델이 실행되도록 하려면 최소한 이런 기본 함수들을 구현해야 해요.

    • init_device: 워커의 장치를 설정하는 함수
    • initialize_cache: 워커의 캐시 설정
    • load_model: 모델 가중치를 장치에 로드
    • get_kv_cache_spec: 모델의 KV 캐시 스펙 생성
    • determine_available_memory: 모델의 피크 메모리 사용량을 프로파일링해 KV 캐시에 쓸 수 있는 메모리(현재 OOM 없이)를 결정
    • initialize_from_config: 지정된 kv_cache_config로 장치 KV 캐시 할당
    • execute_model: 매 스텝마다 호출돼 모델을 추론

    추가로 구현할 수 있는 함수들:

    • 슬립 모드 기능을 지원하려면 sleep과 wakeup 함수 구현
    • 그래프 모드 기능을 지원하려면 compile_or_warm_up_model 함수 구현
    • 스펙큘레이티브 디코딩을 지원하려면 take_draft_token_ids 함수 구현
    • lora를 지원하려면 add_lora, remove_lora, list_loras, pin_lora 함수 구현
    • 데이터 병렬 기능을 지원하려면 execute_dummy_batch 함수 구현

    WorkerBase에 구현 가능한 더 많은 함수는 WorkerBase 클래스를 참고하세요.

  5. my_dummy_attention.py에 어텐션 백엔드 클래스 MyDummyAttention를 구현해요. AttentionBackend를 상속해야 해요. 자기 장치로 어텐션을 계산하는 데 쓰여요. vllm.v1.attention.backends에 많은 구현 예시가 있으니 참고하세요.

  6. 고성능을 위해 커스텀 op를 구현해요. 대부분의 op는 pytorch 네이티브 구현으로도 돌지만 성능은 좋지 않을 수 있어요. 그럴 때 플러그인에 특화된 커스텀 op를 구현할 수 있어요. vLLM이 지원하는 커스텀 op 종류:

    • pytorch ops — 3가지 종류:
      • communicator ops: 장치 커뮤니케이터 op. 예: all-reduce, all-gather. my_dummy_device_communicator.py에 MyDummyDeviceCommunicator 클래스를 구현. DeviceCommunicatorBase를 상속
      • common ops: 일반 op. 예: matmul, softmax. CustomOp 클래스를 참고해 oot 방식으로 등록
      • csrc ops: C++ op. C++로 구현되고 torch 커스텀 op로 등록. csrc 모듈과 vllm._custom_ops를 따라 구현
    • triton ops: triton op는 지금 커스텀 방식이 동작하지 않아요.

  7. (선택) lora, 그래프 백엔드, 양자화, mamba 어텐션 백엔드 등 다른 플러그 가능 모듈도 구현할 수 있어요.

호환성 보장

vLLM은 ModelRegistry.register_model 같은 문서화된 플러그인 인터페이스가 항상 플러그인이 모델을 등록하는 데 사용 가능함을 보장해요. 다만 플러그인이 목표 vLLM 버전과 호환되도록 하는 건 플러그인 개발자의 책임이에요. 예를 들어 "vllm_add_dummy_model.my_llava:MyLlava"는 플러그인이 목표하는 vLLM 버전과 호환되어야 해요.

모델/모듈 인터페이스는 vLLM 개발 중 바뀔 수 있어요. deprecation 로그가 보이면 플러그인을 최신 버전으로 업그레이드하세요.

Deprecation 공지

  • Platform.get_attn_backend_cls의 use_v1 파라미터는 deprecated. v0.13.0에서 제거됨.
  • vllm.attention의 _Backend는 deprecated. v0.13.0에서 제거됨. vllm.v1.attention.backends.registry.register_backend로 AttentionBackendEnum에 새 어텐션 백엔드를 추가하세요.
  • seed_everything 플랫폼 인터페이스는 deprecated. v0.16.0에서 제거됨. vllm.utils.torch_utils.set_random_seed 사용.
  • Platform.validate_request의 prompt는 deprecated. v0.18.0에서 제거됨.

더 알아보기 (Learn more)