vllm-lens 1.0.0__tar.gz

vllm_lens-1.0.0/.github/workflows/checks.yaml

@@ -0,0 +1,45 @@
+name: Checks
+
+on:
+  push:
+    branches:
+      - main
+    paths-ignore:
+      - '.devcontainer/**'
+      - '.github/**'
+      - '.vscode/**'
+      - '.gitignore'
+      - 'README.md'
+  pull_request:
+    branches:
+      - main
+    paths-ignore:
+      - '.devcontainer/**'
+      - '.github/**'
+      - '.vscode/**'
+      - '.gitignore'
+      - 'README.md'
+  workflow_call:
+    inputs:
+      dummy:
+        type: string
+        required: false
+
+jobs:
+  checks:
+    name: Checks
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install uv
+        uses: astral-sh/setup-uv@v4
+      - name: Set up Python
+        run: uv python install
+      - name: Install dependencies
+        run: uv sync
+      - name: Ruff lint
+        run: uv run ruff check
+      - name: Ruff format
+        run: uv run ruff format --check
+      - name: Pyright
+        run: uv run pyright

vllm_lens-1.0.0/.github/workflows/release.yaml

@@ -0,0 +1,29 @@
+name: Release
+
+on:
+  release:
+    types: [published]
+
+jobs:
+  checks:
+    name: Run checks
+    uses: ./.github/workflows/checks.yaml
+
+  release:
+    name: Release to PyPI
+    needs: checks
+    runs-on: ubuntu-latest
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install uv
+        uses: astral-sh/setup-uv@v4
+      - name: Set up Python
+        run: uv python install
+      - name: Set version from tag
+        run: sed -i "s/version=\".*\"/version=\"${GITHUB_REF_NAME#v}\"/" pyproject.toml
+      - name: Build
+        run: uv build
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

vllm_lens-1.0.0/.gitignore

@@ -0,0 +1,30 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv
+
+# Environment variables
+.env
+
+# Logs
+logs/
+scratch/
+
+# Core dump
+**/core
+
+# Scratch files
+**/scratch.ipynb
+
+# Container images (built from containers/*.def)
+*.sif
+
+# Other
+benchmarks.json
+.claude/settings.json

vllm_lens-1.0.0/.pre-commit-config.yaml

@@ -0,0 +1,17 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.15.0
+    hooks:
+      - id: ruff
+        args: [--no-fix]
+      - id: ruff-format
+        args: [--check]
+
+  - repo: local
+    hooks:
+      - id: pyright
+        name: pyright
+        entry: uv run --no-sync pyright
+        language: system
+        types: [python]
+        pass_filenames: false

vllm_lens-1.0.0/.python-version

@@ -0,0 +1 @@
+3.12

vllm_lens-1.0.0/.vscode/extensions.json

@@ -0,0 +1,16 @@
+{
+    "recommendations": [
+        "anthropic.claude-code",
+        "charliermarsh.ruff",
+        "davidanson.vscode-markdownlint",
+        "donjayamanne.githistory",
+        "elagil.pre-commit-helper",
+        "littlefoxteam.vscode-python-test-adapter",
+        "ms-python.python",
+        "ms-python.vscode-pylance",
+        "ms-toolsai.jupyter",
+        "tamasfe.even-better-toml",
+        "ukaisi.inspect-ai",
+        "johnpapa.vscode-peacock",
+    ]
+}

vllm_lens-1.0.0/.vscode/settings.json

@@ -0,0 +1,25 @@
+{
+    "python.envFile": "${workspaceFolder}/.env",
+    "editor.codeActionsOnSave": {
+        "source.organizeImports": "explicit"
+    },
+    "editor.formatOnSave": true,
+    "evenBetterToml.formatter.allowedBlankLines": 1,
+    "evenBetterToml.formatter.arrayAutoCollapse": true,
+    "evenBetterToml.formatter.arrayAutoExpand": true,
+    "evenBetterToml.formatter.arrayTrailingComma": true,
+    "evenBetterToml.formatter.columnWidth": 88,
+    "evenBetterToml.formatter.compactArrays": true,
+    "evenBetterToml.formatter.compactEntries": true,
+    "evenBetterToml.formatter.compactInlineTables": true,
+    "evenBetterToml.formatter.indentEntries": true,
+    "evenBetterToml.formatter.indentString": "    ",
+    "evenBetterToml.formatter.indentTables": true,
+    "evenBetterToml.formatter.inlineTableExpand": true,
+    "evenBetterToml.formatter.reorderArrays": true,
+    "evenBetterToml.formatter.reorderKeys": true,
+    "evenBetterToml.formatter.trailingNewline": true,
+    "evenBetterToml.schema.enabled": true,
+    "evenBetterToml.schema.links": true,
+    "evenBetterToml.syntax.semanticTokens": false,
+}

vllm_lens-1.0.0/PKG-INFO

@@ -0,0 +1,93 @@
+Metadata-Version: 2.4
+Name: vllm-lens
+Version: 1.0.0
+Summary: vLLM plugin for interacting with activations during inference
+License-Expression: MIT
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.12
+Requires-Dist: datasets>=4.0.0
+Requires-Dist: pydantic>=2.0
+Requires-Dist: vllm>=0.16.0
+Requires-Dist: zstandard>=0.23.0
+Description-Content-Type: text/markdown
+
+# vLLM-lens
+
+Extract residual-stream activations and apply steering vectors (including activation oracles) to any vLLM model during inference. Typically hundreds of times faster than HF Transformers, with the ability to generate text, extract activations, and run activation oracles in the same dynamic batch.
+
+Auto-registers as both a [vLLM general plugin](https://docs.vllm.ai/en/latest/design/plugin_system.html) and an [Inspect](https://inspect.aisi.org.uk/) model provider on install. Interact with model internals per-call via `SamplingParams.extra_args` (vLLM) or `GenerateConfig.extra_body` (Inspect).
+
+Supports async, sync, and serve modes with `mp` and `ray` distributed backends for multi-node tensor parallelism. In offline modes, activations are returned directly on `output.activations`. In serve mode (`vllm serve`), they are serialised in the HTTP response and decoded with `decode_activations()`.
+
+## Install
+
+```bash
+uv add vllm-lens
+```
+
+## Examples
+
+These examples use the Inspect integration. See the [`examples/`](examples/) folder for offline and online vLLM usage.
+
+### Inspect AI provider
+
+An [Inspect AI](https://inspect.ai-safety-institute.org.uk/) model provider registered as `vllm-lens`. Extends the built-in vLLM provider to serialize `torch.Tensor` steering vectors for HTTP transport and decode base64-encoded activations from responses into `ModelOutput.metadata["activations"]`. Also supports LoRA adaptors.
+
+Usage is the same as the [default vLLM provider](https://inspect.aisi.org.uk/providers.html#vllm) but with the `vllm-lens` prefix (e.g. `vllm-lens/meta-llama/Llama-3.1-1B`).
+
+#### Extracting activations
+
+```python
+capture_config = GenerateConfig(
+    temperature=0.0,
+    max_tokens=1,
+    extra_body={
+        "extra_args": {"output_residual_stream": extraction_activation_layers},
+        "chat_template_kwargs": {"enable_thinking": False},
+    },
+)
+output = await model.generate(state.messages, config=capture_config)
+residual_stream = output.metadata["activations"]["residual_stream"]
+```
+
+#### Steering with an Activation Oracle
+
+```python
+from vllm_lens import SteeringVector
+
+messages = [ChatMessageUser(content=oracle_content)]
+oracle_config = GenerateConfig(
+    temperature=0.0,
+    max_tokens=50,
+    extra_body={
+        "extra_args": {
+            "apply_steering_vectors": [
+                SteeringVector(
+                    activations=act_vector,
+                    layer_indices=[injection_layer],
+                    scale=steering_coefficient,
+                    norm_match=True,
+                    position_indices=[special_pos],
+                )
+            ],
+        },
+        "lora_request": {
+            "lora_name": "oracle",
+            "lora_int_id": 1,
+            "lora_path": lora_path,
+        },
+        "chat_template_kwargs": {"enable_thinking": False},
+    },
+)
+response = await model.generate(messages, config=oracle_config)
+```
+
+## Theory
+
+Ships a vLLM worker extension that adds hooks to the model in the underlying workers (credit to Sid Black for pioneering this), plus a plugin to track which activations belong to which prompts across vLLM's dynamic batching with paged attention (similar to the nnsight approach).
+
+## Credits
+
+Developed by Alan Cooney, with credit going to Sid Black for the orignal vLLM worker extension idea.

vllm_lens-1.0.0/README.md

@@ -0,0 +1,78 @@
+# vLLM-lens
+
+Extract residual-stream activations and apply steering vectors (including activation oracles) to any vLLM model during inference. Typically hundreds of times faster than HF Transformers, with the ability to generate text, extract activations, and run activation oracles in the same dynamic batch.
+
+Auto-registers as both a [vLLM general plugin](https://docs.vllm.ai/en/latest/design/plugin_system.html) and an [Inspect](https://inspect.aisi.org.uk/) model provider on install. Interact with model internals per-call via `SamplingParams.extra_args` (vLLM) or `GenerateConfig.extra_body` (Inspect).
+
+Supports async, sync, and serve modes with `mp` and `ray` distributed backends for multi-node tensor parallelism. In offline modes, activations are returned directly on `output.activations`. In serve mode (`vllm serve`), they are serialised in the HTTP response and decoded with `decode_activations()`.
+
+## Install
+
+```bash
+uv add vllm-lens
+```
+
+## Examples
+
+These examples use the Inspect integration. See the [`examples/`](examples/) folder for offline and online vLLM usage.
+
+### Inspect AI provider
+
+An [Inspect AI](https://inspect.ai-safety-institute.org.uk/) model provider registered as `vllm-lens`. Extends the built-in vLLM provider to serialize `torch.Tensor` steering vectors for HTTP transport and decode base64-encoded activations from responses into `ModelOutput.metadata["activations"]`. Also supports LoRA adaptors.
+
+Usage is the same as the [default vLLM provider](https://inspect.aisi.org.uk/providers.html#vllm) but with the `vllm-lens` prefix (e.g. `vllm-lens/meta-llama/Llama-3.1-1B`).
+
+#### Extracting activations
+
+```python
+capture_config = GenerateConfig(
+    temperature=0.0,
+    max_tokens=1,
+    extra_body={
+        "extra_args": {"output_residual_stream": extraction_activation_layers},
+        "chat_template_kwargs": {"enable_thinking": False},
+    },
+)
+output = await model.generate(state.messages, config=capture_config)
+residual_stream = output.metadata["activations"]["residual_stream"]
+```
+
+#### Steering with an Activation Oracle
+
+```python
+from vllm_lens import SteeringVector
+
+messages = [ChatMessageUser(content=oracle_content)]
+oracle_config = GenerateConfig(
+    temperature=0.0,
+    max_tokens=50,
+    extra_body={
+        "extra_args": {
+            "apply_steering_vectors": [
+                SteeringVector(
+                    activations=act_vector,
+                    layer_indices=[injection_layer],
+                    scale=steering_coefficient,
+                    norm_match=True,
+                    position_indices=[special_pos],
+                )
+            ],
+        },
+        "lora_request": {
+            "lora_name": "oracle",
+            "lora_int_id": 1,
+            "lora_path": lora_path,
+        },
+        "chat_template_kwargs": {"enable_thinking": False},
+    },
+)
+response = await model.generate(messages, config=oracle_config)
+```
+
+## Theory
+
+Ships a vLLM worker extension that adds hooks to the model in the underlying workers (credit to Sid Black for pioneering this), plus a plugin to track which activations belong to which prompts across vLLM's dynamic batching with paged attention (similar to the nnsight approach).
+
+## Credits
+
+Developed by Alan Cooney, with credit going to Sid Black for the orignal vLLM worker extension idea.

vllm_lens-1.0.0/pyproject.toml

@@ -0,0 +1,88 @@
+[project]
+    classifiers=[
+        "Development Status :: 3 - Alpha",
+        "Intended Audience :: Science/Research",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    ]
+    dependencies=[
+        "datasets>=4.0.0",
+        "pydantic>=2.0",
+        "vllm>=0.16.0",
+        "zstandard>=0.23.0",
+    ]
+    description="vLLM plugin for interacting with activations during inference"
+    license="MIT"
+    name="vllm-lens"
+    readme="README.md"
+    requires-python=">=3.12"
+    version="1.0.0" # Updated by CICD
+
+[build-system]
+    build-backend="hatchling.build"
+    requires=["hatchling"]
+
+[tool.hatch.build.targets.wheel]
+    packages=["vllm_lens"]
+
+[tool.uv]
+    override-dependencies=["aiobotocore>=3.0.0", "transformers>=5.3.0"]
+    [tool.uv.sources]
+        torch=[
+            {index="pytorch-cu126", marker="sys_platform == 'linux' or sys_platform == 'win32'"},
+        ]
+        torchvision=[
+            {index="pytorch-cu126", marker="sys_platform == 'linux' or sys_platform == 'win32'"},
+        ]
+
+    [[tool.uv.index]]
+        explicit=true
+        name="pytorch-cu126"
+        url="https://download.pytorch.org/whl/cu126"
+
+[dependency-groups]
+    dev=[
+        "datasets>=4.6.0",
+        "inspect-ai>=0.3.183",
+        "ipykernel>=7.2.0",
+        "ipywidgets>=8.1.8",
+        "jaxtyping>=0.3.9",
+        "jinja2>=3.1.6",
+        "peft>=0.18.1",
+        "pre-commit>=4.5.1",
+        "pyright>=1.1.408",
+        "pytest-asyncio>=1.3.0",
+        "pytest>=9.0.2",
+        "python-dotenv>=1.2.1",
+        "ray[serve]>=2.54.0",
+        "ruff>=0.15.3",
+        "syrupy>=5.1.0",
+        "torch>=2.9.1",
+        "tqdm>=4.67.3",
+        "transformers>=5.3.0",
+        "typer>=0.24.1",
+    ]
+
+[tool.pytest.ini_options]
+    asyncio_default_test_loop_scope="module"
+    asyncio_mode="auto"
+    filterwarnings=[
+        "ignore::DeprecationWarning:importlib._bootstrap",
+        "ignore:This process.*is multi-threaded.*use of fork:DeprecationWarning",
+    ]
+
+[tool.pyright]
+    include=["vllm_lens"]
+    reportMissingImports=false
+    reportMissingModuleSource=false
+    reportMissingTypeStubs=false
+    stubPath="typestubs"
+    useLibraryCodeForTypes=false
+
+[tool.ruff.lint]
+    ignore=["E402", "F722", "F821"]
+
+    [project.entry-points."vllm.general_plugins"]
+        activations="vllm_lens._activations_plugin:register"
+
+    [project.entry-points.inspect_ai]
+        vllm_lens="vllm_lens._inspect_entry"

vllm_lens-1.0.0/typestubs/transformers/__init__.pyi

@@ -0,0 +1,3 @@
+from typing import Any
+
+def __getattr__(name: str) -> Any: ...