ProbeX: inline safety monitoring for production LLM serving

ProbeX runs a frozen linear safety probe inline on a live OpenAI-compatible server, reading the model's residual stream as it decodes, and measures the three things a team has to settle before putting a probe on a production endpoint: what it costs, whether it survives quantization, and what prefix caching hides from it. Built and measured end-to-end on vLLM and SGLang, Llama-3.1-8B / 3.3-70B / Qwen2.5-7B, on A100/H100 (Modal).

Result

extraction overhead (CUDA graphs ON):   tap ~0%   vs   forced-eager +194% (vLLM) / +76% (SGLang) TPOT @ conc 1
extraction correctness:                  XSTest AUROC 0.994,  cosine-vs-HF 0.9999   (metric-identical on both engines)
quantization robustness:                 FP8 / FP8-KV / AWQ-INT4 / GPTQ-INT4 all benign; crude INT4-RTN breaks the
                                         fragile privacy probe (L20 −0.23); a no-new-labels check catches it (ROC-AUC 1.0)
RadixAttention monitoring blind spot:    last-token probe cache-safe (cos 0.99997); interior monitors 94.4% blind
                                         on a shared-prefix workload; force-recompute (the general fix) costs ~2.2× throughput

	finding
Extraction is free with CUDA graphs on	the standard register_forward_hook tool forces eager and pays +194% (vLLM) / +76% (SGLang) TPOT at low batch; a graph-capturable custom-op tap adds ~0%
The tap matches the offline probe	tap-extracted layer-16 residuals reproduce the offline probe (refusal test AUROC 1.0, XSTest 0.994, cosine-vs-HF 0.9999), metric-identical on both engines with graphs on
Quantization breakage is detectable	FP8 / FP8-KV / AWQ-INT4 / GPTQ-INT4 leave the probe intact; only crude INT4-RTN breaks the fragile privacy probe (L20 ΔAUROC −0.23, CI excludes 0), and a one-scalar check needing no new labels at serving time flags it: ROC-AUC 1.0, FPR 0 over 53 cells (validated on the same grid that motivated it)
Prefix caching creates a monitoring blind spot	a cached prompt recomputes only its last token → last-token probes are cache-safe at zero cost, but interior-token monitors go 94.4% blind on a shared-prefix workload; force-recompute costs ~2.2× throughput
Generalizes across models	the pattern holds across Llama-3.1-8B, Llama-3.3-70B, Qwen2.5-7B

Full narrative in WRITEUP.md; the cross-engine result in results/sglang/CROSS_ENGINE.md. A follow-on fuses the activation gather and probe matvec into one kernel (Triton and CUDA vs the torch baseline): a tie at decode sizes, 4.7x at large batch, cuBLAS wins multi-probe; see results/kernel/FINDINGS.md.

How it works

• Reading activations from a live server is what costs, because the standard extraction path turns CUDA graphs off. The reference tool (vLLM-Lens) reads activations with a PyTorch forward hook, which doesn't fire under CUDA-graph capture, so it forces enforce_eager=True engine-wide, killing graphs and inflating decode TPOT exactly at the low batch sizes where interactive monitoring runs.
• The fix is a graph-capturable custom op. torch.ops.probex.probex_tap is spliced into the decoder layer and recorded inside the CUDA graph: it copies the layer-L residual into a fixed-address buffer with static-shape GPU ops, no host callback. Graphs stay on.
• The read point transfers across engines. Both vLLM and SGLang use fused add+RMSNorm and return (hidden_states, residual); their sum at layer L equals HF output_hidden_states[L+1], the same read point on both engines.
• Quantization failure is caught with no new serving-time labels by relative calibration retention (cosine of the frozen probe direction vs a fresh difference-of-means, normalized to bf16), which collapses to ~0 exactly on the breaking cells.
• RadixAttention only recomputes the last token of a cached prompt, which is a last-token probe's read point, so it stays covered; interior-token monitors don't.

Reproduce

The headline numbers replay from committed results/*.json, CPU-only, no GPU, no accounts, stdlib-only, ~5 seconds:

python3 -m scripts.reproduce_key_numbers

The measurements regenerate on Modal (HF/Anthropic tokens live in Modal Secrets in your own workspace, never in .env):

cp .env.example .env && modal profile activate <your-workspace>      # then: modal secret list

modal run deploy/app.py        --mode bench        # serving curve + forced-eager A/B (vLLM)
modal run deploy/kernel.py     --mode verify       # graph-capturable tap, correctness gate (vLLM)
modal run deploy/sglang_port.py --mode verify      # same gate on SGLang
modal run deploy/sglang_port.py --mode decomp      # tap vs graphs-on vs eager decomposition
modal run deploy/sglang_port.py --mode radix       # RadixAttention read-point economics
modal run deploy/sglang_port.py --mode act2        # quantization matrix

Each stage writes a result blob and skips if present (idempotent / resumable). Absolute serving latencies are host-sensitive; the ratios (tap ≈ graphs-on, forced-eager penalty, cache-hit fraction) reproduce.

Drop in your own probe

The tap layer is a pluggable config (src/probe_config.py): point PROBEX_PROBE_CONFIG / PROBEX_TAP_LAYER at your probe and the full serving-cost + quantization pipeline runs unchanged.

Layout

Path	Purpose
src/	Engine-agnostic harness: probe.py, extract.py, loadtest.py (multi-process load client), data + stats; the taps probex_kernel.py (vLLM) / sglang_kernel.py (SGLang) and launchers
deploy/	Modal entrypoints: app.py (serving), kernel.py (vLLM tap), sglang_port.py (SGLang port), cascade.py, quantization cells, scores.py
results/	CPU-replayable measurement blobs: vllm/, sglang/ (+ CROSS_ENGINE.md), quantization + cascade artifacts
scripts/reproduce_key_numbers.py	Recompute the headline numbers from results/, CPU-only
WRITEUP.md	The full narrative
tests/	Unit tests for the harness (metric layer is dependency-free)

Scope

ProbeX re-implements published probe methodology (Anthropic/DeepMind cascades; Duan's monitor-staleness study) on open models and extends it into white space: the wall-clock serving dimension, a graph-capturable extraction kernel, and the cross-engine RadixAttention analysis. It does not benchmark any frontier lab's closed, deployed classifier.