Project Summary

We present vLLM Hook, a modular plug-in that extracts and adapts the internal states of vLLM models. This allows for the full support of inference-time monitoring and control methods.

What is vLLM Hook?

Despite achieving high inference throughput, the current vLLM package offers limited support accessing or modifying important internal states of a vLLM model. Examples of these internal states include attentions, attention heads, token embeddings, and activations. The lack of programmability in significant internal states creates a technological gap that prevents the use of advanced techniques for monitoring, adjusting, and controlling deployed models at inference time. To bridge this gap, our vLLM Hook is an open-source plug-in of vLLM to enable flexible programming of internal states within vLLM models.

Our vLLM Hook provides seamless and non-intrusive integration to vLLM. The system overview of vLLM Hook is illustrated in Figure 1. Based on a configuration file (Config) specifying which internal states within a vLLM model to “hook”, vLLM Hook supports two essential features: passive programming and active programming.

  • Passive programming probes and saves the selected internal states for subsequent analysis, while keeping the model generation intact.
  • Active programming enables efficient intervention of model generation by altering the selected internal states.

vLLM Hook overview

Figure 1: Overview of vLLM Hook. With a customizable configuration file (Config), vLLM Hook supports two essential features: active programming and passive programming.

Why do we need vLLM Hook?

We position the vLLM Hook project as a developer kit that supports various use cases requiring the programming of vLLM model internal states. To demonstrate the programmability of different internal states through the vLLM Hook, we developed three examples that address the practical need to capture or modify specific parts of a vLLM model, including attentions, attention heads, and activations.

  • In-model Monitoring [passive programming]: Use vLLM Hook to capture critical “inference traces” generated during inference. The saved states are then used for model monitoring and misalignment detection. One example is the Attention Tracker, which uses selected attention statistics from a transformer-based vLLM model to detect prompt injection attempts. The “in-model safety guardrail” design, enabled by vLLM Hook, differs from the typical “cascaded safety guardrail” design. The latter requires the use of an additional moderation model to inspect the input and output of the deployed model.

    Usage example of Attention Tracker: Link

  • Model Steering [active programming]: Use vLLM Hook to modify the internal states of the model to steer its output toward desired behaviors. Although recent model steering methods demonstrate promising inference-time alignment capabilities without retraining the deployed model (see AI Steerability 360 Toolkit), the current vLLM package prevents most model steering methods that modify internal states from being implemented. One example is Activation Steering, which injects a steering vector into some layer of a deployed model for intervention, such as improving the capability in instruction following.

    Usage example of Activation Steering for Improving Instruction Following: Link

  • Selective Retrieval [passive programming]: Use vLLM Hook to activate a selected set of components within a deployed model for efficient and effective data retrieval. The rationale is to retrieve data embeddings from a deployed model that are relevant to the task of interest. One example is to only activate a selected set of attention heads within a deployed model to improve the re-ranking performance in information retrieval.

    Usage example of Selective Retrieval for Improving Data Re-ranking: Link

Core Functions of vLLM Hook

vLLM-Hook is built as a modular plugin library that enables users to capture internal model signals (passive programming), intervene during inference (active programming), and optionally analyze the captured signals. We give an illustration of vLLM-Hook orchestration in Figure 2. The core of our framework consists of two abstractions: worker and analyzer. On a higher level, worker defines when and how we program within the vLLM pipeline, and analyzer defines the evaluation metrics based on the captured signals (if any). Together with worker and analyzer, a Config File is provided to specify the more granular behavior of the worker and analyzer. For instance, it could include the profile of the important layers and attention heads of a model, the steering vector of a desired model behavior, and the statistic reduction method (take average vs. weighted summation), etc.

vLLM Hook orchestration

Figure 2: An illustration of vLLM Hook orchestration. vLLM Hook is a lightweight wrapper around a native vLLM system and uses llm=HookLLM to initialize an LLM instance, llm.generate for response generation (model forward pass), and llm.analyze for optional in-model analysis.

Below is an example of the code snippet to use the HookLLM class in vLLM Hook to run the Attention Tracker application. We also provide configurations for different model serving and data storage options. Please see configs.md for details.

Offline (HookLLM)

from vllm_hook_plugins import HookLLM
from vllm import SamplingParams

llm = HookLLM(
    model="ibm-granite/granite-3.1-8b-instruct",
    worker_name="probe_hook_qk",
    analyzer_name="attn_tracker",
    config_file="model_configs/attention_tracker/granite-3.1-8b-instruct.json",
)

# rpc (in-memory) path:
out   = llm.generate(text, SamplingParams(...), save_to_disk=False)
stats = llm.analyze(probes=out[0].probes, analyzer_spec={...})

# disk path (artifact under /dev/shm/vllm_hook/<run_id>/):
out   = llm.generate(text, SamplingParams(...), save_to_disk=True, run_id="run-1")
stats = llm.analyze(analyzer_spec={...})  # uses the last run_id

# activation steering (worker_name="steer_hook_act", no analyzer): no save_to_disk, difference is observed by comparing against a use_hook=False baseline.
out_steered = llm.generate(text, SamplingParams(...))
out_plain   = llm.generate(text, SamplingParams(...), use_hook=False)

Format/save-mode are env-vars on the offline driver process, set before HookLLM(...) is constructed (the worker subprocess inherits them at spawn):

VLLM_HOOK_USE_SAFETENSORS=1   # write .safetensors instead of .pt
VLLM_HOOK_ASYNC_SAVE=1        # background daemon thread instead of inline
VLLM_HOOK_USE_SHM=1           # legacy shared-memory fast path (hidden states + last_token only)

Performance Comparison to vLLM Eagle

This section benchmarks vLLM Hook v.s. Native vLLM Eagle (ExampleHiddenStatesConnector) for hidden state extraction. Both systems extract intermediate transformer hidden states during inference for subsequent analysis. The key differences are:

  • vLLM Hook adds a configuration layer to specify which hidden state to extract, whereas vLLM Eagle builds on speculative decoding and requires using an additional draft model to capture the hidden state.
  • For passive programming, vLLM Hook supports the extraction of other hidden features such as attention, which is not supported by vLLM Eagle.
  • For active programming, vLLM Hook supports activation steering, which is not supported by vLLM Eagle.

To compare the efficiency of vLLM Hook v.s. vLLM Eagle, we benchmark their latency and memory usage in extracting token embeddings when varying the number of layers.

What Have We Compared

  • Number of layers: 1 to all
  • Length of prompts: around 16, 64, 256, 512 tokens
  • Extraction modes:
System Token Mode Description
vLLM-Hook last_token Extracts the last-token hidden state per prompt per layer
vLLM-Hook all_tokens Extracts hidden states for all tokens per prompt per layer
Native vLLM Eagle N/A (effectively all_token) Built-in extraction via ExampleHiddenStatesConnector

Key Findings

Figure 3 compares four metrics across 4 prompt lengths (rows) and 28 layer counts (x-axis) using the Qwen2-1.5B-Instruct model and one GPU (A100/H100 class, ~74 GB VRAM).

Benchmark grid: vLLM Hook vs native vLLM Eagle

Figure 3: Comparison between vLLM Hook and vLLM Eagle. Four metrics across 4 prompt lengths (rows) and 28 layer counts (x-axis). Shaded bands show ± standard deviation over 10 runs.

The key findings are summarized as follows:

1. At low layer counts, all three systems perform comparably

When extracting hidden states from only a few layers (~1–4), all three systems show similar generation latency regardless of prompt length. For example, at 1 layer, vLLM Hook (last_token) = 29.8 ms, vLLM Hook (all_tokens) = 32.4 ms, and native vLLM Eagle = 29.2 ms (16-token prompt).

2. vLLM-Hook (last_token) has the lowest and flattest latency overhead

vLLM Hook (last_token) adds only 1.3–2.1× overhead from 1 to 28 layers, independent of prompt length. At 512 tokens and 28 layers it takes 58.4 ms — compared to 103.5 ms for vLLM Hook (all_tokens) and 539.2 ms for native Eagle. This makes it the best choice when only the final-position representation is needed. For example, when one needs hidden states from all layers, we have:

Prompt length vLLM Hook (last_token) vLLM Hook (all_tokens) native vLLM Eagle
16 tokens 39.9 ms 37.9 ms 31.5 ms
64 tokens 47.1 ms 46.3 ms 49.7 ms
256 tokens 45.6 ms 62.3 ms 123.4 ms
512 tokens 58.4 ms 103.5 ms 539.2 ms

vLLM Hook (last_token) slices to the last token at capture time inside the worker, so only a single (hidden_size,) vector per layer is written to disk, regardless of prompt length. Native vLLM stores all prompt tokens and all layers together.

3. Artifact size for last_token is prompt-length invariant

vLLM Hook (last_token) produces a flat 677 KB artifact regardless of prompt length, since it stores only one vector per prompt per layer. This is up to ~510× smaller than all_tokens or native Eagle at long sequences (~344 MB at 28 layers / 512 tokens), making it suitable for high-throughput scenarios where storage or I/O bandwidth is a bottleneck.

4. Native vLLM Eagle carries a ~2.3 GiB GPU memory overhead

Native vLLM ExampleHiddenStatesConnector is built on Eagle-3 speculative decoding infrastructure, which requires loading a dummy drafter model. In practice this reduces the available KV cache by ~2.3 GiB (vLLM-Hook: 18.45 GiB available; native: 16.13 GiB available).

To reproduce the results, please see numerical_analysis.

Summary

We envision vLLM Hook as an organic innovation engine that will enable the development and deployment of advanced techniques for monitoring and adjusting vLLM models. These techniques require the use of programming internal states during inference. We invite the community to contribute to the vLLM Hook project to expand its core functions, use cases, and demonstrations. Significant contributors will be invited to co-author future vLLM technical reports.

References