# ablated-vllm-plugin

A vLLM plugin for per-layer ablation studies.

Dynamically generates ablated subclasses of vLLM models to support ablation of attention and MLP layers of _any vLLM model_.

## Install

To install an editable version of this plugin to your environment (and add new models), first clone the repo:
```bash
git clone https://github.com/maxzuo/ablated-vllm-plugin
```

and install with pip:
```bash
cd ablated-vllm-plugin
pip install -e .
```

To use directly with `pip`:
```bash
pip install ablated-vllm
```

## Usage

To run an ablated model, start by downloading a huggingface model to a local directory:
```bash
hf download Qwen/Qwen3.5-9B --local-dir ./ablated-qwen3.5-9b
```

Then, set `ablated_attention_layers` and/or `ablated_mlp_layers` in the model's `config.json`, and add the `"Ablated"` prefix to the model architecture.
Values are dicts mapping layer index to an implementation:

```json
{
  "architectures": ["AblatedQwen3_5ForConditionalGeneration"],
  "ablated_attention_layers": {"10": "zero"},
  "ablated_mlp_layers": {"5": "identity"}
}
```

### Ablation implementations

| Value | Effect |
|---|---|
| `'zero'` | Replaces the layer output with zeros |
| `'identity'` | Passes hidden states through unchanged (skips the layer) |

### Inference

No changes are needed to run this model with vLLM.

### Adding a new model

Add an entry to `ablated_vllm/models.json`:

```json
{
  "MyModelForCausalLM": {
    "attention_classes": ["MyAttentionClass"],
    "mlp_classes": ["MyMLPClass"],
    "decoder_class": "MyDecoderLayer",
    "module_path": "vllm.model_executor.models.my_model"
  }
}
```

No other code changes needed.
The plugin generates `AblatedMyModelForCausalLM` dynamically and automatically registers it with vLLM.
