Command Palette
Search for a command to run...
vLLM Hook v0: vLLM의 프로그래밍 모델 내부 구조를 위한 플러그인
vLLM Hook v0: vLLM의 프로그래밍 모델 내부 구조를 위한 플러그인
Ching-Yun Ko Pin-Yu Chen
vLLM 을 사용하여 Gemma-3-27B-IT 배포하기
초록
현대 인공지능(AI) 모델은 특히 트랜스포머 기반 대규모 언어 모델(LLM)의 경우, 런타임 효율성과 자원 할당을 최적화하기 위해 추론 엔진에 배포됩니다. vLLM 프로젝트는 모델 서빙과 추론을 지원하기 위한 주요 오픈소스 라이브러리입니다. 그러나 현재 vLLM의 구현은 배포된 모델의 내부 상태 프로그래밍 가능성을 제한합니다. 이로 인해 테스트 타임 모델 정렬 및 향상 방법과 같은 일반적인 기법들의 활용이 불가능해집니다. 예를 들어, 어텐션 패턴을 기반으로 한 적대적 프롬프트 감지나 활성화 스티어링(activation steering)에 기반한 모델 응답 조정이 불가능해집니다. 이러한 중요한 격차를 해소하기 위해, 우리는 vLLM 모델의 내부 상태 프로그래밍을 가능하게 하는 오픈소스 플러그인인 vLLM Hook을 제시합니다. vLLM Hook은 캡처할 내부 상태를 지정하는 구성 파일을 기반으로 vLLM과 원활하게 통합되며, 수동 프로그래밍(passive programming)과 능동 프로그래밍(active programming)이라는 두 가지 핵심 기능을 지원합니다. 수동 프로그래밍의 경우, vLLM Hook은 모델 생성을 그대로 유지한 채 선택된 내부 상태를 프로빙(probing)하여 후속 분석을 가능하게 합니다. 능동 프로그래밍의 경우, vLLM Hook은 선택된 내부 상태를 변경함으로써 모델 생성에 대한 효율적인 개입을 가능하게 합니다.
One-sentence Summary
The authors present vLLM Hook, an open-source plug-in for the vLLM inference engine that enables configurable programming of internal model states through passive probing and active intervention, thereby overcoming existing programmability constraints to support test-time alignment, adversarial prompt detection, and activation steering for large language models.
Key Contributions
- vLLM Hook is an open-source plugin that enables configuration-driven programming of internal states within the vLLM inference engine, directly addressing the limitation that restricts test-time model alignment and enhancement methods.
- The system implements two core programming modes, passive programming for non-intrusive state probing that preserves generation, and active programming for real-time intervention via the alteration of selected internal states.
- Three practical demonstrations validate the plugin, showcasing prompt injection detection, enhanced retrieval-augmented retrieval, and activation steering to verify its utility for runtime model monitoring and adjustment.
Introduction
Modern large language models rely on inference engines like vLLM to optimize deployment efficiency and resource allocation. However, the current vLLM implementation restricts access to and modification of internal model states during inference, which blocks essential test-time alignment techniques such as adversarial prompt detection and activation steering. To address this limitation, the authors develop vLLM Hook, an open-source plug-in that enables precise programming of internal states through a simple configuration file. The framework supports passive probing for real-time analysis and active intervention to directly alter model outputs, effectively unlocking practical applications like enhanced retrieval-augmented generation and secure prompt monitoring.
Dataset
- The authors do not provide a dataset description in the submitted text.
- Dataset composition and sources: The content only outlines a GitHub contribution workflow and references a repository URL. No data sources or composition details are included.
- Key details for each subset: The text contains no information regarding subset sizes, origins, or filtering criteria.
- How the paper uses the data: No training splits, mixture ratios, or data processing steps are described.
- Cropping strategy, metadata construction, or other processing details: None are mentioned in the provided material.
Method
The vLLM-Hook framework is designed as a modular plugin system that enables both passive and active programming within the vLLM inference pipeline. At its core, the framework operates through two primary abstractions: the worker and the analyzer, which are orchestrated by a configuration file that defines the behavior of each component. The worker integrates directly into the vLLM runtime and is responsible for either capturing internal model states during inference (passive programming) or modifying the model's behavior in real time (active programming). This integration is achieved by subclassing the standard vLLM GPU worker and overriding the load_model method to install PyTorch forward hooks on selected model modules. These hooks are applied to specific attention layers and heads, as specified in the configuration, allowing for targeted observation or intervention during the forward pass.
As shown in the figure below, the framework begins with a native vLLM system that receives an input prompt. The user specifies the components to probe via a configuration file, which is then used to guide the vLLM-Hook system. The system captures internal states during inference, which can be either saved for later analysis or used to enable active programming, such as model steering or customized generation. The configuration file defines the model identity, important layers and attention heads, and the mode of signal capture—such as whether to collect data for all tokens or only the last token. These configurations are managed through a lightweight registry and a HookLLM wrapper class that initializes the LLM instance and interfaces with the core vLLM engine.
The workflow proceeds in three stages: configuration identification, probing, and programming. In the configuration stage, the user identifies the components to probe, potentially using external data. During probing, the worker measures targeted model internals via hooks during inference, capturing relevant activations or attention weights. The final stage involves programming, where the saved states are used either for passive monitoring—such as evaluating prompt injection risks—or for active intervention, such as steering model behavior. This process is illustrated in the framework diagram, where the vLLM-Hook plugin wraps the vLLM system and interacts with the LLMEngine, which manages input processing, scheduling, model execution, and output processing.
The analyzer component operates on the saved states after inference completion. It retrieves the cached data using a unique run identifier and reassembles the desired statistics, such as attention weights, to compute metrics like prompt injection attack scores or document relevance scores. This is achieved through a modular analyzer class that takes the hook directory and layer-to-head mappings as inputs and processes the cached data to compute specific metrics. The analyzer is triggered via the llm.analyze method, which allows users to perform post-inference analysis without modifying the core model or runtime. This modular design enables the framework to support a wide range of applications, including safety monitoring, model steering, and selective retrieval, by combining different workers and analyzers within the same orchestration system.