← Writing

Running Safety Probes on Production Inference Servers

systems safety interpretability inference

TL;DR.

Epistemic status. Llama-3.1-8B, 3.3-70B, and Qwen2.5-7B on A100/H100 via Modal, across two production engines (vLLM 0.23, SGLang 0.5.14). Numbers reproduce from cached artifacts in about 5 seconds (scripts/reproduce_key_numbers.py). The extraction op is instrumentation, not a base-engine speedup, and the privacy hard set is small and templated. Code and full writeup.

What inline monitoring actually costs

There is a tidy story in the interpretability literature: train a linear probe on the residual stream, get AUROC 0.95 on “is this harmful,” ship it as a cheap safety monitor. The probe’s arithmetic is free, a dot product against a hidden state the model already computed, so people quote the FLOPs ratio (probe FLOPs over forward-pass FLOPs, roughly 0) and call it solved.

But the real cost is getting the hidden state out of the serving engine.

The reference open-source tool for this (BEIS’s vLLM-Lens) reads activations with a register_forward_hook that clones the layer output. Forward hooks do not fire when the model runs under torch.compile or inside a captured CUDA graph, so the moment you load the plugin it sets enforce_eager=True at engine init and turns CUDA graphs off for the entire engine, whether or not any individual request asked for activations. Decode is launch-bound, CUDA graphs are the thing that hides per-kernel launch latency, so killing them is expensive precisely at low batch:

Llama-3.1-8B, A100, 1k/1k prompts, concurrency 1, forced-eager vs graphs-on TPOT: +194% on vLLM, +76% on SGLang.

Low batch is the interactive, latency-sensitive regime, which is where you would run a real-time safety monitor. So the “probes are free” claim holds only for the arithmetic. The standard way to feed the probe doubles to triples your per-token latency at low batch.

The fix: a graph-capturable custom op

The fix is to stop using a host-side Python callback and instead give the graph something it can capture. I register a real custom op, torch.ops.probex.probex_tap, and splice it into the decoder layer’s forward pass. A registered op is an opaque node in the graph: torch.compile will not trace into it and will not refuse to capture it, and CUDA-graph capture records it like any other kernel. Inside, it does only static-shape GPU work: gather the layer-L residual into a fixed, pre-allocated staging buffer, no host callback, no synchronization. CUDA graphs stay on.

Two things made the port to a second engine work.

The read point is the same on both. vLLM and SGLang both implement Llama with fused add+RMSNorm, where a decoder layer returns (hidden_states, residual) and the true residual stream after block L is hidden_states + residual, which is exactly HF’s output_hidden_states[L+1]. So “tap the sum of the two returns at layer L” is one sentence that works identically on both. Tap-extracted layer-16 residuals reproduce the offline probe: refusal test AUROC 1.0, XSTest AUROC 0.994, cosine 0.9999 against the HF reference, with graphs on. The two engines’ AUROCs match to five decimals. That is metric-identical extraction; I did not diff the tensors bitwise.

The hard part was different on each. On vLLM the wrinkle was the V1 architecture running the model in a separate EngineCore subprocess: the op has to be installed inside that process before graph capture, via a vllm.general_plugins entry point. On SGLang the wrinkle was torch.compile: SGLang 0.5 captures the prefill graph with a piecewise-compiled backend, and my first version read per-request layout tensors inside the compiled region, which tripped a Dynamo shape-guard and crashed with “runtime recompilation.” SGLang’s own code only touches that layout in the eager glue between compiled submodules, so the fix was to make the in-graph op layout-free (copy the whole step’s residual) and move the per-request gather out to the eager wrapper. At decode, where the latency penalty actually lives, copying the whole step is the per-request gather, so it costs the same.

Result, measured as a three-way decomposition (tap / graphs-on-no-tap / forced-eager, same launcher so the only variable is the tap):

The tap adds ~0% TPOT vs graphs-on at every concurrency on both engines, versus the +194% and +76% the forced-eager tool pays.

The op itself is a memory-bound index_select and copy (arithmetic intensity roughly 0). I estimate its cost at about 12 µs on A100, roughly 0.12% of a decode step, from device specs and CUDA-graph-replay timing; the Nsight profile that would confirm it was blocked on Modal’s managed hosts, so treat that figure as an estimate.

To be clear about what this is: an instrumentation kernel. The base engine runs no faster. It just removes a tooling tax, so leaving the safety monitor on in production no longer costs latency.

Does the monitor survive quantization?

Production does not serve bf16. So I froze the probe on bf16 and swapped the served model underneath it across the formats real deployments ship: FP8 weights, FP8 KV-cache, calibrated INT4 (AWQ and GPTQ), plus a crude INT4-RTN as a worst-case control.

The calibrated formats are benign: frozen-probe XSTest AUROC moves within noise (FP8-KV Δ −0.0001 to −0.0003; AWQ −0.0039; GPTQ −0.0013), and this holds cross-engine and cross-scale (the same pattern on Llama-3.3-70B and Qwen2.5-7B). That extends Duan’s quantization-robustness result to FP8 and to the live serving setting.

The interesting failure is the crude control. The privacy probe flags requests for a named individual’s private PII against benign requests for public contact information; its hard split is about 120 templated prompts where both classes name a person, separable only by public-versus-private, which makes it the most fragile target in Duan’s study. Uncalibrated INT4-RTN drops that probe’s hard-split AUROC at deep layers (L20 ΔAUROC −0.232, bootstrap CI [−0.299, −0.169], excludes zero), while the refusal probe shrugs it off. At serving time you have no labels to notice the monitor degraded, so the useful contribution here is the detector. It is a single scalar, relative calibration retention: the cosine between the frozen bf16 probe direction and a fresh difference-of-means direction, computed on the quantized model’s activations over the same labeled eval set and normalized to the bf16 value. It collapses to about 0 on the breaking cells and stays put on the benign ones. Over 53 quant cells it gives ROC-AUC 1.0 on the deep breakages and zero false positives, needing no new labels at serving time. That grid is also the one that motivated the detector, so the 1.0 is an in-grid consistency check rather than held-out generalization.

The part that only exists in the serving layer

This part does not show up in a notebook. SGLang’s RadixAttention caches the KV of shared prefixes, so a repeated or shared-prefix prompt is never recomputed, which means an inline monitor reading the residual stream at prefill has nothing to read for the cached tokens. I expected this to be a wall for inline monitoring. What I measured turned out to be more specific.

I measured what actually gets recomputed on a cache hit. SGLang caps the prefix match at length − 1: it always recomputes at least the final token, because it needs that position’s output to produce the next one. And the final token is the read point of a last-token probe. So:

So the design choice is simple. Reading the residual back out of the KV cache is infeasible: the cache stores K/V projections, not the layer-L residual stream, so you cannot reconstruct it. Force-recomputing every monitored request restores full coverage but forfeits the cache speedup, about 2.2× throughput at concurrency 64. Extracting only on cache-miss tokens is free and complete for a last-token probe (the last token is always a miss) and has the 94% hole for interior probes. That gives a simple rule: a last-token design is RadixAttention-safe at zero extra cost, and anything reading interior tokens has to pay the recompute tax or accept the blind spot. This constraint on how you build production monitors only shows up when you look at the monitor and the cache together.

Why I built it this way

ProbeX is a systems project built around an interpretability question. The probe itself is the easy part: I could have written the classifier in an afternoon, and the literature already has. I wanted to know what happens when you actually run one on a production endpoint, because that is where the interesting costs are, and almost none of them are visible from the offline activation-extraction scripts most of this work uses.

The thing I am least sure about is how much of the RadixAttention result survives a different caching policy, and that is where I would look first if I picked this back up.

The whole thing reproduces from committed JSON in about five seconds on a laptop, and the GPU runs are one modal run each. Code and the full writeup are here.