fix: let a static pod declare the layer a readout axis reads

apps/inference/README.md

@@ -92,6 +92,16 @@ attention and DFA. Because the set is fixed at load, this is a routing decision:
 a pod that did not declare a point cannot grow one, and says so in a 400 that
 lists what it did declare.
 
+`STATIC_POINTS_EXTRA` adds sites to a set the server resolves for itself —
+`sae` or `sae+auto` — as a JSON list in the same spelling, `["resid_post.40"]`.
+It is for readout axes, which are the one thing a pod cannot resolve at startup:
+an axis is a database row that arrives with the request, so the layer it reads is
+not knowable when the graphs are recorded. `sae` declares the SAE sites and
+nothing else, and the alternatives are both wrong for one extra layer — `auto`
+declares every layer, which the larger pods cannot afford, and an explicit list
+declares no writes, which trades every steer for the readout. Declared as a read
+and a write, so an axis can also be steered at the layer it was fitted at.
+
 `GENERATION_ONLY=true` is the empty case, worth up to +249% decode on a 1B model
 and ~1% at 8B — small-model completion traffic and nothing else. The server
 refuses to start if `SAE_SETS` is non-empty (an SAE read *is* a capture), if

apps/inference/neuronpedia_inference/args.py

@@ -67,6 +67,17 @@ def parse_env_and_args():
     # Declare attn at a layer to keep attention/DFA. Anything not declared 400s, so this is a
     # routing decision, not a tuning knob.
     args.static_points = os.getenv("STATIC_POINTS")
+    # Sites to declare on top of a set this server resolves for itself, as a JSON list in the same
+    # spelling STATIC_POINTS takes: '["resid_post.40"]' or '[["resid_post", 40]]'.
+    #
+    # It exists for readout axes, which are the one thing a pod cannot resolve at startup. An axis
+    # is a database row that arrives with the request, so which layers this pod will be asked to
+    # capture is not knowable when the graphs are recorded -- and a layer that was not declared
+    # then cannot be added to a running static pod at all. STATIC_POINTS=sae covers the SAE sites
+    # and nothing else, and a pod whose axes read elsewhere has no other way to say so: `auto`
+    # would declare every layer, which several pods cannot afford, and an explicit list declares
+    # no writes, which trades every steer for the readout.
+    args.static_points_extra = os.getenv("STATIC_POINTS_EXTRA")
     # Renamed in interp-engine 1.3 along with the backend it drives. Refused rather than aliased,
     # because a deploy that sets the old name is asking for a tap set and would otherwise get a
     # hooked pod: same endpoints, several times slower, and no line in the log saying why.

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

@@ -24,7 +24,12 @@
     features_to_vllm_steering_spec,
     resolve_max_new_tokens,
 )
-from neuronpedia_inference.engine_adapter import BackendUnsupported, assert_steering_available, get_tokenize
+from neuronpedia_inference.engine_adapter import (
+    BackendUnsupported,
+    assert_capture_layers_declared,
+    assert_steering_available,
+    get_tokenize,
+)
 from neuronpedia_inference.inference_utils.persona import (
     AxisAsset,
     AxisRequestError,
@@ -178,6 +183,18 @@
             logger.warning(f"Axis request rejected: {exc}", exc_info=True)
             return JSONResponse(content={"error": str(exc)}, status_code=exc.status_code)
 
+        # Asked before anything is rendered or generated. A static pod's tap set is fixed when its
+        # graphs are recorded, and an axis names its layer at request time, so this is the one
+        # mismatch that no startup check could have caught.
+        try:
+            assert_capture_layers_declared(
+                model,
+                [axis.layer for axis in axes],
+                f"Readout axis {[axis.id for axis in axes]}",
+            )
+        except BackendUnsupported as exc:
+            return JSONResponse(content={"error": str(exc)}, status_code=400)
+
     # Every requested axis has to agree about how the conversation is rendered, because those
     # conditions are applied before generation and so change the text itself. There is no way to
     # render one conversation two ways in a single generation, and honouring one axis's

apps/inference/neuronpedia_inference/engine_adapter.py

@@ -189,6 +189,43 @@ def assert_residual_available(model: object, what: str = "This endpoint", point:
     )
 
 
+def assert_capture_layers_declared(
+    model: object,
+    layers: list[int],
+    what: str = "This endpoint",
+    point: str = "resid_post",
+) -> None:
+    """Refuse a capture at a layer a static pod did not declare, by layer and not just by name.
+
+    :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.
+
+    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.
+    """
+    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:
+        return
+    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]}\"]')."
+    )
+
+
 def _vllm_points_use_native_resid(model: object, points: list[_CapturePoint]) -> bool:
     """True when these resid_post reads should go through native extract, not static/hooks."""
     if not points or not all(p.address.name == "resid_post" and p.address.layer is not None for p in points):

apps/inference/neuronpedia_inference/server.py

@@ -22,6 +22,7 @@
     check_cuda_driver,
     load_model,
     select_backend,
+    to_address,
 )
 
 from neuronpedia_inference.args import parse_env_and_args
@@ -245,6 +246,54 @@ def _parse_static_points(raw: str | None) -> Any:
     return points
 
 
+def _parse_extra_static_points(raw: str | None) -> list[Address]:
+    """``STATIC_POINTS_EXTRA``: a JSON list of addresses to declare beside a resolved set.
+
+    Same spelling as an explicit ``STATIC_POINTS`` list, so a site moves between the two by
+    cut-and-paste. Empty and unset both mean nothing extra, which is the common case.
+    """
+    if raw is None or raw.strip() == "":
+        return []
+    try:
+        points = json.loads(raw)
+    except json.JSONDecodeError as exc:
+        raise ValueError(
+            f"STATIC_POINTS_EXTRA={raw!r} is not JSON (e.g. '[\"resid_post.40\"]' or '[[\"resid_post\", 40]]')."
+        ) from exc
+    if not isinstance(points, list):
+        raise ValueError(f"STATIC_POINTS_EXTRA={raw!r} must be a JSON list of addresses.")
+    return [to_address(point) for point in points]
+
+
+def _with_extra_points(
+    reads: list[Address],
+    writes: list[Address],
+    extra: list[Address],
+) -> tuple[list[Address], list[Address]]:
+    """Add ``STATIC_POINTS_EXTRA`` to a resolved set, as reads and as writes.
+
+    Both, for the reason ``auto`` implies its writes: a capture site that cannot be written is a
+    readout that works and then refuses every steer at the same layer, and steering on a persona
+    direction at the layer it was fitted at is a thing this server is asked to do. One layer both
+    ways costs two buffers, which is the cheap half of the trade `sae+auto` gets wrong at scale.
+
+    Deduplicated, so naming a site an SAE already covers is free rather than doubled, and appended
+    after the resolved set so the startup log reads as "the SAE set, plus what was declared on top".
+    """
+    read_keys = {str(address) for address in reads}
+    write_keys = {str(address) for address in writes}
+    merged_reads = list(reads)
+    merged_writes = list(writes)
+    for address in extra:
+        if str(address) not in read_keys:
+            merged_reads.append(address)
+            read_keys.add(str(address))
+        if str(address) not in write_keys:
+            merged_writes.append(address)
+            write_keys.add(str(address))
+    return merged_reads, merged_writes
+
+
 def _with_residual_set(
     model: Any,
     reads: list[Address],
@@ -718,6 +767,18 @@ def load_model_and_sae():
                 f"STATIC_POINTS={static_mode} needs SAE_SETS so there are hook sites to declare. "
                 "Use STATIC_POINTS=auto for the residual set alone."
             )
+        extra_points = _parse_extra_static_points(getattr(args, "static_points_extra", None))
+        if extra_points and static_mode not in _SAE_RESOLVED_MODES:
+            # Refused rather than ignored. Every other value already says everything it is going
+            # to: `auto` declares each layer, an explicit list is the caller's own set, and unset
+            # is a hooked pod where nothing needs declaring. Merging into one of those would be a
+            # no-op that reads, in a deploy config, as though a site had been added.
+            raise ValueError(
+                f"STATIC_POINTS_EXTRA has nothing to add to STATIC_POINTS={static_mode!r}. It "
+                "declares sites beside a set this server resolves after the SAEs load, so it "
+                "applies to STATIC_POINTS=sae or sae+auto. Add the sites to the list itself, or "
+                "drop STATIC_POINTS_EXTRA."
+            )
         num_gpus = max(1, int(getattr(args, "num_gpus", 1) or 1))
         if num_gpus > 1:
             logger.info("Multi-GPU: sharding across %d GPUs (%s)", num_gpus, args.backend)
@@ -872,9 +933,12 @@ def load_model_and_sae():
             reads, writes = sae_static_addresses(SAEManager._instance)
             if static_mode == "sae+auto":
                 reads, writes = _with_residual_set(model, reads, writes, num_layers)
+            if extra_points:
+                reads, writes = _with_extra_points(reads, writes, extra_points)
             logger.info(
-                "STATIC_POINTS=%s: freezing %d read site(s) and %d write site(s)",
+                "STATIC_POINTS=%s%s: freezing %d read site(s) and %d write site(s)",
                 static_mode,
+                f" + {[str(a) for a in extra_points]}" if extra_points else "",
                 len(reads),
                 len(writes),
             )

apps/inference/tests/unit/test_generation_only.py

@@ -110,6 +110,35 @@ def test_steered_generation_needs_declared_writes_not_just_reads(self):
             )
         )
 
+    def test_an_axis_layer_the_pod_did_not_declare_is_refused_by_layer_not_by_point_name(self):
+        """The 70B case: `resid_post` IS declared, at the layer its SAE reads and not the one the
+        axis was fitted at. Checked by address, or the pod passes here and dies inside generate."""
+        from interp_engine.address import Address
+
+        from neuronpedia_inference.engine_adapter import assert_capture_layers_declared
+
+        pod = SimpleNamespace(
+            hooks_available=False,
+            graph_replay=True,
+            static_points=(Address("resid_post", 50),),
+            static_writes=(),
+        )
+        with pytest.raises(BackendUnsupported) as excinfo:
+            assert_capture_layers_declared(pod, [40], "Readout axis ['lu_assistant-axis']")
+        message = str(excinfo.value)
+        assert "lu_assistant-axis" in message
+        assert "[40]" in message
+        # Names the way out, since this is a relaunch and not a code change.
+        assert "STATIC_POINTS_EXTRA" in message
+
+        assert_capture_layers_declared(pod, [50], "Readout axis ['x']")
+
+    def test_a_hooked_pod_declares_nothing_and_can_capture_anywhere(self):
+        from neuronpedia_inference.engine_adapter import assert_capture_layers_declared
+
+        assert_capture_layers_declared(SimpleNamespace(), [40])
+        assert_capture_layers_declared(SimpleNamespace(hooks_available=True), [40])
+
     def test_an_unknown_backend_is_assumed_capable(self):
         # Anything without the attribute predates it and hooks in-process (eager, and the test
         # doubles for it). Defaulting to "refuse" here would break every such caller.

apps/inference/tests/unit/test_vllm_backend_kwargs.py

@@ -19,13 +19,16 @@
 from __future__ import annotations
 
 import pytest
+from interp_engine import Address
 
 from neuronpedia_inference.server import (
     _DEFAULT_CUDAGRAPH_CAPTURE_SIZES,
     _engine_context_len,
+    _parse_extra_static_points,
     _parse_static_points,
     _vllm_backend_kwargs,
     _vllm_engine_backend,
+    _with_extra_points,
 )
 
 
@@ -145,6 +148,42 @@ def test_a_malformed_static_points_is_refused_rather_than_ignored():
         _parse_static_points("resid_post")
 
 
+def test_extra_static_points_takes_both_address_spellings_and_means_nothing_when_unset():
+    assert _parse_extra_static_points(None) == []
+    assert _parse_extra_static_points("") == []
+    assert _parse_extra_static_points("  ") == []
+    assert [str(a) for a in _parse_extra_static_points('["resid_post.40"]')] == ["resid_post.40"]
+    assert [str(a) for a in _parse_extra_static_points('[["resid_post", 40]]')] == ["resid_post.40"]
+
+
+def test_a_malformed_extra_static_points_is_refused_rather_than_ignored():
+    with pytest.raises(ValueError, match="STATIC_POINTS_EXTRA"):
+        _parse_extra_static_points("resid_post.40")
+    with pytest.raises(ValueError, match="STATIC_POINTS_EXTRA"):
+        _parse_extra_static_points('"resid_post.40"')
+
+
+def test_an_extra_point_is_declared_to_read_and_to_write():
+    """Both, so a persona axis can be steered at the layer it was fitted at and not only read."""
+    reads, writes = _with_extra_points(
+        [Address("resid_post", 50)],
+        [Address("resid_post", 49)],
+        [Address("resid_post", 40)],
+    )
+    assert [str(a) for a in reads] == ["resid_post.50", "resid_post.40"]
+    assert [str(a) for a in writes] == ["resid_post.49", "resid_post.40"]
+
+
+def test_an_extra_point_the_sae_set_already_covers_costs_nothing_twice():
+    reads, writes = _with_extra_points(
+        [Address("resid_post", 40)],
+        [Address("resid_post", 40)],
+        [Address("resid_post", 40), Address("resid_post", 40)],
+    )
+    assert [str(a) for a in reads] == ["resid_post.40"]
+    assert [str(a) for a in writes] == ["resid_post.40"]
+
+
 def test_no_endpoint_accepts_an_image():
     """The premise the override rests on, checked rather than asserted in a comment.