fix: make a static pod's tap set visible, and refuse a bad steer with a 400

apps/inference/neuronpedia_inference/config.py

@@ -130,6 +130,10 @@ class Config:
     sae_gpu_budget_gib: str | None = None
     # Cap on page-locked host memory for those master copies. None = measure it.
     sae_pinned_host_gib: float | None = None
+    # Human-readable summary of what STATIC_POINTS (+ STATIC_POINTS_EXTRA) asked to declare, for
+    # the startup banner. A string rather than the addresses because this banner prints before the
+    # SAEs load, so `sae` has no address list yet; the resolved set is in the ready banner.
+    static_points_summary: str | None = None
 
     # Accepted as constructor arguments but stored under the derived names below.
     include_sae: InitVar[list[str] | None] = None
@@ -169,6 +173,7 @@ class Config:
         "exclude_sae_patterns",
         "backend",
         "num_gpus",
+        "static_points_summary",
     )
 
     def __post_init__(

apps/inference/neuronpedia_inference/endpoints/steer/completion.py

@@ -1,5 +1,5 @@
 import logging
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Sequence
 from typing import Any, cast
 
 import torch
@@ -20,7 +20,13 @@
 from interp_engine import steer as engine_steer
 
 from neuronpedia_inference.config import Config
-from neuronpedia_inference.engine_adapter import BackendUnsupported, assert_steering_available, tlens_hook_to_point
+from neuronpedia_inference.engine_adapter import (
+    BackendUnsupported,
+    assert_steer_layers_declared,
+    assert_steering_available,
+    declares_static_taps,
+    tlens_hook_to_point,
+)
 from neuronpedia_inference.inference_utils.steering import (
     SteeringSettings,
     format_sse_message,
@@ -137,6 +143,14 @@
             status_code=400,
         )
 
+    # Asked here because the spec that writes is built inside the generator, where a refusal is a
+    # 500 mid-stream rather than a status this can return.
+    if NPSteerType.STEERED in request.types and declares_static_taps(model):
+        try:
+            assert_steer_layers_declared(model, steer_write_layers(features))
+        except BackendUnsupported as exc:
+            return JSONResponse(content={"error": str(exc)}, status_code=400)
+
     max_new_tokens, no_room = resolve_max_new_tokens(len(tokens), int(request.n_completion_tokens))
     if no_room is not None:
         return no_room
@@ -261,6 +275,38 @@
     )
 
 
+def steer_layer_for_hook(hook_name: str) -> int:
+    """Which layer a steer at ``hook_name`` writes: ``resid_pre[X]`` lands at ``X-1``.
+
+    Split out so the pre-flight write check and the spec that does the writing read the layer the
+    same way. Computing it twice is how a pod passes a check for one layer and fails the engine at
+    another.
+    """
+    if "resid_post" in hook_name:
+        return int(hook_name.split(".")[1])
+    if "resid_pre" in hook_name:
+        return int(hook_name.split(".")[1]) - 1
+    raise ValueError(f"Unsupported hook name for vLLM steering: {hook_name}")
+
+
+def steer_write_layers(features: Sequence[NPSteerFeature | NPSteerVector]) -> list[int]:
+    """The layers a steer over ``features`` will write to, sorted and deduplicated.
+
+    Answerable in the request handler, before a StreamingResponse has taken the reply away, where
+    :func:`features_to_vllm_steering_spec` runs inside the generator and cannot return a status.
+    Hooks it cannot map are left to that function, whose ValueError is already handled.
+    """
+    sae_manager = SAEManager.get_instance()
+    layers: set[int] = set()
+    for feature in features:
+        hook_name = sae_manager.get_sae_hook(feature.source) if isinstance(feature, NPSteerFeature) else feature.hook
+        try:
+            layers.add(steer_layer_for_hook(hook_name))
+        except ValueError:
+            continue
+    return sorted(layers)
+
+
 def features_to_vllm_steering_spec(settings: SteeringSettings) -> SteeringSpec:
     """Build an engine ``SteeringSpec`` (per-layer Add/ProjectionCap ops) for the vLLM backend.
 
@@ -271,12 +317,7 @@
     layer_features: dict[int, list[tuple[NPSteerFeature | NPSteerVector, torch.Tensor]]] = {}
     for feature in settings.features:
         hook_name = sae_manager.get_sae_hook(feature.source) if isinstance(feature, NPSteerFeature) else feature.hook
-        if "resid_post" in hook_name:
-            layer = int(hook_name.split(".")[1])
-        elif "resid_pre" in hook_name:
-            layer = int(hook_name.split(".")[1]) - 1
-        else:
-            raise ValueError(f"Unsupported hook name for vLLM steering: {hook_name}")
+        layer = steer_layer_for_hook(hook_name)
 
         steering_vector = torch.tensor(feature.steering_vector, dtype=torch.float32)
         if not torch.isfinite(steering_vector).all():

apps/inference/neuronpedia_inference/endpoints/steer/completion_chat.py

@@ -23,11 +23,14 @@
     _feature_to_steerspec,
     features_to_vllm_steering_spec,
     resolve_max_new_tokens,
+    steer_write_layers,
 )
 from neuronpedia_inference.engine_adapter import (
     BackendUnsupported,
     assert_capture_layers_declared,
+    assert_steer_layers_declared,
     assert_steering_available,
+    declares_static_taps,
     get_tokenize,
 )
 from neuronpedia_inference.inference_utils.persona import (
@@ -297,6 +300,16 @@
             status_code=400,
         )
 
+    # The write-side twin of the axis-layer check above, and asked for the same reason: the spec
+    # that writes is built inside the SSE generator, so a refusal there is a 500 mid-stream. A
+    # projection cap writes wherever the vector it caps was fitted, which on this endpoint is not
+    # the layer the axis reads -- the 70B pod declared 40 for the readout and was asked to write 32.
+    if wants_steering and declares_static_taps(model):
+        try:
+            assert_steer_layers_declared(model, steer_write_layers(features))
+        except BackendUnsupported as exc:
+            return JSONResponse(content={"error": str(exc)}, status_code=400)
+
     # Convert promptChatFormatted to NPSteerChatMessage for the readouts, so they analyze the
     # same conversation (including the system message, blanked or not) that generation saw.
     inputPromptForAxes = [NPSteerChatMessage(role=msg["role"], content=msg["content"]) for msg in promptChatFormatted]

apps/inference/neuronpedia_inference/engine_adapter.py

@@ -189,6 +189,73 @@ def assert_residual_available(model: object, what: str = "This endpoint", point:
     )
 
 
+def declares_static_taps(model: object) -> bool:
+    """True when this pod's tap set is fixed, so a per-layer check is worth asking at all.
+
+    Guards the work of resolving which layers a request would touch, not just the answer: on a
+    hooked pod that resolution reaches into the SAE manager for each feature's hook name, and a
+    check that cannot refuse anything should not be able to raise on the way to saying so.
+    """
+    return not getattr(model, "hooks_available", True)
+
+
+def _resid_aliases(address: Address) -> tuple[Address, ...]:
+    """The other name for the same residual add: ``resid_pre[L]`` IS ``resid_post[L-1]``.
+
+    Mirrors ``resid_stream_aliases`` in the engine's vllm_backend, which is what the refusals
+    these two checks pre-empt match against. Copied rather than imported: it is three lines, and
+    the import would reach past the package root into the vLLM backend module.
+
+    Drift is the risk that trades for, so these checks lean PERMISSIVE. A site this fails to
+    recognize as declared is one the engine refuses for itself, which is the 500 they exist to
+    replace and not a new failure; a site they wrongly call missing would refuse a request the pod
+    could have served.
+    """
+    if address.layer is None:
+        return (address,)
+    layer = int(address.layer)
+    if address.name == "resid_pre" and layer > 0:
+        return (address, Address("resid_post", layer - 1))
+    if address.name == "resid_post":
+        return (address, Address("resid_pre", layer + 1))
+    return (address,)
+
+
+def _undeclared_layers(declared: set[str], layers: list[int], point: str) -> list[int]:
+    """Which of ``layers`` no alias of ``point`` covers, sorted and deduplicated."""
+    return sorted(
+        {
+            layer
+            for layer in layers
+            if not any(str(alias) in declared for alias in _resid_aliases(Address(point, layer)))
+        }
+    )
+
+
+def _static_layer_miss(
+    model: object,
+    layers: list[int],
+    point: str,
+    *,
+    private_attr: str,
+    public_attr: str,
+) -> tuple[list[int], list[str]] | None:
+    """``(missing layers, declared)`` for a static pod, or None when the question does not apply.
+
+    None covers both "hooks run, so every site is reachable" and "nothing declared at all", the
+    second because that pod is refused wholesale by :func:`assert_hooks_available` /
+    :func:`assert_steering_available` and a layer list would be the wrong reason to give.
+    """
+    if getattr(model, "hooks_available", True):
+        return None
+    declared = getattr(model, private_attr, None) or getattr(model, public_attr, ()) or ()
+    declared_keys = {str(address) for address in declared}
+    if not declared_keys:
+        return None
+    missing = _undeclared_layers(declared_keys, layers, point)
+    return (missing, sorted(declared_keys)) if missing else None
+
+
 def assert_capture_layers_declared(
     model: object,
     layers: list[int],
@@ -200,29 +267,53 @@ def assert_capture_layers_declared(
     :func:`assert_residual_available` asks whether ``point`` is declared at all, which is the
     right question for an endpoint that reads every layer or none. A readout axis reads exactly
     one, so a pod can declare ``resid_post`` and still not have the layer in hand: the 70B pod
-    declares the site its layer-50 SAE reads and is asked for layer 40 by an axis fitted there.
+    declared the site its layer-50 SAE reads and was asked for layer 40 by an axis fitted there.
 
     That distinction only started to matter when axes became database rows. A shipped asset was
     on disk before the graphs were recorded, so the mismatch was a deploy-time fact; an axis that
     arrives with the request makes it a per-request one, and the engine's own refusal comes from
-    inside the generate call -- a 500 with a traceback, after the prompt was rendered. Asked here
-    instead, it is a 400 that names the axis and the layer, which is a pod that needs
-    ``STATIC_POINTS_EXTRA`` rather than a bug.
+    inside the generate call -- a 500 with a traceback, after the prompt was rendered.
     """
-    if getattr(model, "hooks_available", True):
-        return
-    declared = getattr(model, "_static_reads", None) or getattr(model, "static_points", ()) or ()
-    declared_keys = {str(address) for address in declared}
-    if not declared_keys:
-        return  # Not a static pod; the hooks/native checks own this case.
-    missing = sorted({layer for layer in layers if f"{point}.{layer}" not in declared_keys})
-    if not missing:
+    miss = _static_layer_miss(model, layers, point, private_attr="_static_reads", public_attr="static_points")
+    if miss is None:
         return
+    missing, declared = miss
     raise BackendUnsupported(
         f"{what} reads {point} at layer(s) {missing}, which this pod did not declare. Its CUDA "
         f"graphs are recorded against a fixed tap set, so no layer can be added while it runs. "
-        f"Declared: {sorted(declared_keys)}. Relaunch it with those layers in "
-        f"STATIC_POINTS_EXTRA (e.g. STATIC_POINTS_EXTRA='[\"{point}.{missing[0]}\"]')."
+        f"Declared: {declared}. Relaunch it with those layers declared: STATIC_POINTS=auto covers "
+        f"every layer, and STATIC_POINTS_EXTRA adds a closed set beside STATIC_POINTS=sae "
+        f"(e.g. STATIC_POINTS_EXTRA='[\"{point}.{missing[0]}\"]')."
+    )
+
+
+def assert_steer_layers_declared(
+    model: object,
+    layers: list[int],
+    what: str = "Steered generation",
+    point: str = "resid_post",
+) -> None:
+    """The write-side twin of :func:`assert_capture_layers_declared`.
+
+    :func:`assert_steering_available` asks whether this pod declared ANY write site, which catches
+    a generation-only pod and misses the case that actually bit: writes declared, at other layers.
+    A projection cap writes wherever the feature or vector it caps was fitted, and an uploaded
+    vector can name any layer at all, so the set is not one a pod can enumerate at startup --
+    which is why the pod this was written for now declares every layer instead.
+
+    Kept separate from the read check because the two sets genuinely differ: an SAE set declares
+    writes under the ``resid_pre[L]`` -> ``resid_post[L-1]`` mapping, so a pod can hold a read at
+    a layer and no write there.
+    """
+    miss = _static_layer_miss(model, layers, point, private_attr="_static_writes", public_attr="static_writes")
+    if miss is None:
+        return
+    missing, declared = miss
+    raise BackendUnsupported(
+        f"{what} writes {point} at layer(s) {missing}, which this pod did not declare. Its CUDA "
+        f"graphs are recorded against a fixed tap set, so no write tap can be added while it "
+        f"runs. Declared for writing: {declared}. Relaunch it with STATIC_POINTS=auto, which "
+        f"declares every layer to read and to write."
     )