Per-request activation steering for Gemma 4 31B with vLLM

I wanted activation steering that I could change between messages without giving up vLLM throughput. steering-gemma is the resulting system: a Textual client, an OpenAI-compatible FastAPI server, a vLLM inference worker on one H100, and a separate Hugging Face training job for hidden-state extraction.

A paired H100 benchmark measured 42.8 tok/s with steering and 42.7 tok/s with zero steering. That 0.1 tok/s difference is within run-to-run noise. The model applies a real residual-stream intervention on 30 layers without paying a measurable decode penalty.

Three problems took most of the work. The vectors had to remain effective on a heavily instruction-tuned model, mutable request state had to survive Dynamo and CUDA graph capture, and metadata requests had to remain responsive during an H100 cold start.

Residual-stream injection across layers 20–49

The actual intervention is small. Each Gemma decoder block updates the residual stream passed to the next block; the patch adds the weighted sum of the active vectors to that update:

$$h'_{\ell} = h_{\ell} + \operatorname{block}_{\ell}(h_{\ell}) + \sum_{j=1}^{k} \alpha_j d_{j,\ell}$$

(1)

In Equation 1, $d(j, ℓ)$ is vector $j$ at layer $ℓ$, and $α_j$ is its signed magnitude. A request specifies the magnitudes. The vectors remain fixed during inference.

Gemma 4 31B has 60 decoder blocks. I steer layers 20–49, the range that behaved best in development testing, and write zeros to the remaining layers. Active vectors are summed before decoding, so combining five personas still costs one addition per selected layer.

Contrastive mean-difference vector training

Each vector is defined by a positive and negative persona description. The built-in pirate vector contrasts “a swashbuckling pirate captain from the 1700s, full of nautical jargon” with “a stuffy university professor of literature, full of academic jargon.” Both descriptions are paired with the same 34 neutral suffixes.

The chat template mattered much more than I expected. The trainer places each suffix in an assistant turn and captures the final non-padding token from every selected layer. With raw prompt concatenation, the pirate vector was weak even at large magnitudes. Capturing a state from persona-conditioned assistant text made the behavior obvious around magnitude 4.

For each layer, the trainer subtracts the negative activation from its paired positive activation, averages the differences, and normalizes the result:

$$d_{\ell} = \frac{\frac{1}{N}\sum_{i=1}^{N}\left(h^{+}_{i,\ell} - h^{-}_{i,\ell}\right)}{\left\lVert\frac{1}{N}\sum_{i=1}^{N}\left(h^{+}_{i,\ell} - h^{-}_{i,\ell}\right)\right\rVert_2}$$

(2)

I also expected PCA to win. It did not. The leading component emphasized variation among suffix topics: weather, food, work, dreams. Mean difference retained the activation shift shared across those suffixes and produced stronger persona effects at lower magnitudes.

Effective magnitudes were usually between ±3 and ±6. Smaller values were often overridden by instruction tuning. Values beyond ±8 frequently reduced coherence.

Development outputs at magnitude ±5

The following samples use the request “tell me about your morning in two sentences.” They are development examples; no controlled evaluation was run.

The verbose -5 run returned one period. It is still my favorite result from the project: the negative direction reduced an otherwise normal answer to the smallest possible response.

Vectors can be composed linearly. happy +5, pirate +3, and poetic +2 produced “I woke up with a giant smile and danced across the sandy shores as the sun popped peek-a-boo!” The negative formal direction, formal -5, produced “Not much! Just chilling in the cloud. ☁️ What's up with you?”

These are empirical activation directions, and they can drift across prompts, sampling settings, or model versions. The outputs are compelling, but this experiment does not establish a universal semantic axis.

Failed contrast: tsundere versus deretsun

A vector trained from “a tsundere anime girl” and “a deretsun anime girl (not tsundere)” had unit norm across all 30 selected layers and produced almost nothing at magnitude 8. I initially trusted the norm and assumed the training code was fine. The norm was the wrong thing to trust.

The negative description was probably too obscure. If Gemma has a weak representation for deretsun, both descriptions primarily activate the broader “anime girl” concept. Normalization then gives the residual noise unit norm and makes it look healthy. A useful contrast needs two concepts that the base model represents clearly.

Mutable steering state inside compiled vLLM execution

This section consumed most of the implementation time. Hugging Face eager inference with register_forward_hook worked immediately and measured 9.7 tok/s. The compatible torch.compile path reached 10.4 tok/s after a long cold compilation. vLLM supplied the throughput I wanted, then removed the Python-level hook behavior the project depended on.

I tried four state mechanisms. ContextVar.get() caused a graph break. A module global was specialized to its trace-time value. An import-time tensor pointed at stale storage after vLLM initialized its worker process. The fourth attempt used register_buffer, a standard PyTorch primitive that happened to give Dynamo exactly the ownership semantics this project needed.

The patch replaces Gemma4DecoderLayer.__init__ before vLLM constructs the model. Each layer registers a zero-filled, non-persistent steering_direction buffer, and the patched forward adds it on every call. Zero steering still executes the addition; it simply adds zeros.

The hot path consists of in-place tensor writes. The worker combines the requested vectors, writes each layer buffer with copy_(), generates the response, and zeroes the buffers in a finally block. The storage address never changes, so the captured graph reads new values without recompiling. persistent=False keeps the runtime buffer out of model checkpoints while preserving normal module device and dtype handling.

CPU API and scale-to-zero GPU deployment

The first backend placed FastAPI, vector metadata, and inference on one H100 class. It was simple and unpleasant to use. After scale-down, opening the TUI could not even list available vectors until the 31B model had loaded and compiled.

The deployed version separates the API, inference, and training processes:

The CPU Server keeps one container warm and reads vector metadata from a Modal Volume. The H100 Engine starts on the first chat call and scales down after 15 idle minutes. Vector training runs as a separate H100 function because it requires Hugging Face hidden-state extraction.

A full inference cold start still takes several minutes for weight loading, kernel compilation, and CUDA graph capture. That delay now starts when the user sends a chat request. Health checks, sessions, and vector metadata remain available throughout it.

Cold-start and request-queue failures

One cold start sat at 0% weight loading for twenty minutes. Modal mounts Volumes through 9P, and vLLM did not classify that mount as a network filesystem, so it disabled parallel safetensor prefetching. Setting safetensors_load_strategy="prefetch" explicitly fixed the stall.

The readiness check produced a second failure. It polled /v1/chat/completions, and each timed-out probe remained queued on an engine configured with max_inputs=1. Once startup completed, the engine processed seven canceled probes before the real request. Readiness polling now uses /health on the CPU service and never enters the GPU queue.

TUI controls and per-message experiment state

The Textual client exposes one signed magnitude control per vector. Each card includes half-step and full-step adjustments plus a zero reset. Zero is the only off state; adding a separate boolean toggle would create two sources of truth. Positive and negative values select opposite poles of the trained contrast.

The chat pane streams Markdown and reports time to first token, decode throughput, token count, and interruption state. The message schema includes vector names and magnitudes, which allows a turn to be replayed or debugged later. Chat history lives in local SQLite; vector files remain on the Modal Volume.

Limitations and planned experiments

Mid-generation magnitude changes are the interaction I want next. The current API fixes steering state for one generation; changing it during decoding would require resuming from the existing KV-cache prefix after updating the buffers.

Mixture-of-experts models are the other unresolved case. Residual directions may interact with expert routing, and useful control may require per-expert vectors or interventions around router outputs. I have not tested that path yet.

The repository includes deployment instructions, vector training commands, API smoke tests, and the Textual client: tetraslam/steering-gemma.

Sources