feat(inference): --quantization and --kv_cache_dtype; bump interp-engine to 1.8.0

apps/graph/pyproject.toml

@@ -12,7 +12,7 @@ dependencies = [
   "pydantic>=2.11.4",
   "torch>=2.11.0",
   "transformers>=4.57.1",
-  "interp-engine==1.7.3",
+  "interp-engine==1.8.0",
   "python-dotenv>=1.1.0,<2",
   "hf-transfer>=0.1.9,<0.2",
   "psutil>=7.0,<8",

apps/inference/README.md

@@ -72,6 +72,18 @@ error, after the KV cache size already looked fine.
 `VLLM_GPU_MEMORY_UTILIZATION` env var overrides the flag (as env vars do for
 every `start.py` argument).
 
+**Load precision:** `--quantization fp8` quantizes a bf16 checkpoint at load, with
+no calibration step: every linear layer narrows to one byte and the embeddings
+stay at `--model_dtype`, so Llama-3.3-70B goes from ~131 GiB to ~68 GiB and fits
+one 96 GB card. `fp8` is vLLM-only; `bnb-4bit` works on both backends and
+`bnb-8bit` on eager. `--kv_cache_dtype fp8` halves vLLM's KV cache instead, which
+about doubles the context that fits. Both reach `interp_engine.load_model` by the
+same names (interp-engine >= 1.8) and a scheme the resolved backend cannot apply is
+refused at startup. A checkpoint that already ships quantized needs neither flag.
+Env `MODEL_QUANTIZATION` / `KV_CACHE_DTYPE` override the flags. Size the result
+with the engine's `gpu-sizer` before launching:
+`python gpu-sizer/fit.py meta-llama/Llama-3.3-70B-Instruct --quantization fp8`.
+
 **Which vLLM backend a pod runs:** interp-engine splits vLLM into three, and two
 env vars pick between them. They differ only in what the pod declares up front,
 which is what decides whether it can keep CUDA graphs.

apps/inference/neuronpedia_inference/args.py

@@ -21,6 +21,11 @@ def parse_env_and_args():
     # "auto" loads each checkpoint in its native dtype (e.g. gemma-2/3, qwen3 = bf16; gpt2 = fp32)
     # rather than forcing fp32. Override with MODEL_DTYPE=float32 for tight fp32 numerics if needed.
     args.model_dtype = os.getenv("MODEL_DTYPE", "auto")
+    # On-load quantization of a wider checkpoint ("fp8", "bnb-4bit", "bnb-8bit"), and the dtype of
+    # vLLM's paged KV cache ("auto" or "fp8"). Both go to load_model by the same names; the engine
+    # refuses a scheme the resolved backend cannot apply. Unset = load the checkpoint as stored.
+    args.quantization = os.getenv("MODEL_QUANTIZATION", "") or ""
+    args.kv_cache_dtype = os.getenv("KV_CACHE_DTYPE", "auto") or "auto"
     args.sae_dtype = os.getenv("SAE_DTYPE", "float32")
     args.token_limit = int(os.getenv("TOKEN_LIMIT", "200"))
     # Separate cap for the lens endpoints only (logit/jacobian lens). Defaults to

apps/inference/neuronpedia_inference/config.py

@@ -98,6 +98,11 @@ class Config:
     custom_hf_model_id: str | None = None
     sae_sets: list[str] = field(default_factory=lambda: ["res-jb"])
     model_dtype: str = "auto"
+    # On-load quantization scheme ("" = the checkpoint as stored) and the vLLM KV cache dtype
+    # ("auto" = model_dtype). Passed to load_model by the same names; the KV dtype also sizes
+    # the serving limits, since an fp8 cache holds twice the tokens.
+    quantization: str = ""
+    kv_cache_dtype: str = "auto"
     sae_dtype: str = "float32"
     secret: str | None = None
     port: int = 5000
@@ -160,6 +165,8 @@ class Config:
         "custom_hf_model_id",
         "override_model_id",
         "model_dtype",
+        "quantization",
+        "kv_cache_dtype",
         "sae_dtype",
         "port",
         "token_limit",

apps/inference/neuronpedia_inference/server.py

@@ -1,5 +1,6 @@
 import asyncio
 import gc
+import inspect
 import json
 import logging
 import os
@@ -391,6 +392,39 @@ def _vllm_backend_kwargs(
     return kwargs
 
 
+def _load_precision_kwargs(quantization: str, kv_cache_dtype: str) -> dict[str, Any]:
+    """The ``load_model`` arguments for on-load quantization and the KV cache dtype, when set.
+
+    Only set values are passed, so a pod that asks for neither reaches ``load_model`` with the same
+    call as before. Both names arrived in interp-engine 1.8: on an older engine (a venv synced before
+    the pin moved) they would fall into ``**backend_kwargs`` and die in a constructor as an opaque
+    ``TypeError``, so the signature is checked here and the refusal says what to do.
+    """
+    kwargs: dict[str, Any] = {}
+    if quantization:
+        kwargs["quantization"] = quantization
+    if kv_cache_dtype not in ("auto", ""):
+        kwargs["kv_cache_dtype"] = kv_cache_dtype
+    if not kwargs:
+        return kwargs
+    accepted = inspect.signature(load_model).parameters
+    missing = [name for name in kwargs if name not in accepted]
+    if missing:
+        raise ValueError(
+            f"{', '.join(missing)} needs interp-engine >= 1.8, and this venv's load_model does not take "
+            f"it. Sync the venv to the pin in pyproject.toml (`uv sync`), or drop the flag."
+        )
+    return kwargs
+
+
+def _kv_cache_dtype_for_sizing(model_dtype: str, kv_cache_dtype: str) -> str:
+    """The dtype the KV cache is really held in, for the per-token cost the serving limits use."""
+    if kv_cache_dtype in ("auto", ""):
+        return model_dtype
+    # vLLM spells its 8-bit caches fp8 / fp8_e4m3 / fp8_e5m2; startup_memory prices "float8".
+    return "float8" if kv_cache_dtype.startswith("fp8") else kv_cache_dtype
+
+
 def _resolve_generation_only(args: Any) -> bool:
     """Validate ``GENERATION_ONLY`` against the rest of the startup config, and refuse early.
 
@@ -567,7 +601,10 @@ def _log_ready_banner(elapsed_seconds: float) -> None:
         [
             f"model: {model_desc}",
             f"backend: {config.backend} | device: {config.device} | gpus: {config.num_gpus} | "
-            f"model dtype: {config.model_dtype} | sae dtype: {config.sae_dtype}",
+            f"model dtype: {config.model_dtype}"
+            + (f" | quantization: {config.quantization}" if config.quantization else "")
+            + (f" | kv cache dtype: {config.kv_cache_dtype}" if config.kv_cache_dtype != "auto" else "")
+            + f" | sae dtype: {config.sae_dtype}",
             f"saes: {len(sae_manager.loaded_saes)} resident of {configured_saes} configured "
             f"({', '.join(sae_manager.valid_sae_sets) or 'none'})",
             f"token limits: prompt={config.token_limit} activation={config.activation_token_limit} "
@@ -924,6 +961,8 @@ def load_model_and_sae():
             custom_hf_model_id=custom_hf_model_id,
             sae_sets=args_sae_sets,
             model_dtype=args.model_dtype,
+            quantization=args.quantization,
+            kv_cache_dtype=args.kv_cache_dtype,
             sae_dtype=args.sae_dtype,
             token_limit=args.token_limit,
             lens_token_limit=args.lens_token_limit,
@@ -991,12 +1030,19 @@ def load_model_and_sae():
                 "default_prepend_bos": True,
                 "model_kwargs": config.model_kwargs,
             }
+        # On-load quantization and the KV cache dtype, by the names load_model takes. The engine
+        # narrows the linear layers (embeddings stay at dtype) and refuses a scheme this backend
+        # cannot apply, so a wrong flag fails here rather than as an OOM later.
+        precision_kwargs = _load_precision_kwargs(config.quantization, config.kv_cache_dtype)
+        if precision_kwargs:
+            logger.info("Load precision: %s", precision_kwargs)
         model = load_model(
             hf_model_id,
             backend=engine_backend,
             device=args.device,
             dtype=config.model_dtype,
             num_gpus=num_gpus,
+            **precision_kwargs,
             **backend_kwargs,
         )
 
@@ -1008,20 +1054,21 @@ def load_model_and_sae():
         # the per-request token budget. On vLLM we admit up to max_concurrent (vLLM
         # batches); off vLLM (eager) we serve one at a time. See startup_memory.py.
         is_vllm = isinstance(model, VLLMModel)
+        kv_dtype = _kv_cache_dtype_for_sizing(config.model_dtype, config.kv_cache_dtype)
         if is_vllm:
             attn = model._attn_dims  # type: ignore[attr-defined]
             model_info = ModelMemoryInfo(
                 n_layers=num_layers,
                 n_kv_heads=attn["n_kv_heads"],
                 head_dim=attn["head_dim"],
-                dtype=config.model_dtype,
+                dtype=kv_dtype,
             )
         else:
             model_info = ModelMemoryInfo(
                 n_layers=num_layers,
                 n_kv_heads=model.n_kv_heads,  # type: ignore[attr-defined]
                 head_dim=model.head_dim,  # type: ignore[attr-defined]
-                dtype=config.model_dtype,
+                dtype=kv_dtype,
             )
         serving_limits = compute_serving_limits(device=args.device, is_vllm=is_vllm, model_info=model_info)
         config.set_max_tokens(serving_limits.max_tokens)

apps/inference/pyproject.toml

@@ -52,7 +52,10 @@ dependencies = [
   # 1.7.3 tags a system turn the chat template injects on its own (Llama's date
   # preamble) as `role=system` with no message index, so the lens and NLA chats can
   # show it as its own bubble. 1.7.1-1.7.2 are capture-side fixes.
-  "interp-engine[vllm,quant]==1.7.3",
+  # 1.8.0 adds `quantization=` and `kv_cache_dtype=` to load_model, which start.py's
+  # --quantization / --kv_cache_dtype pass through: FP8 on load is what fits Llama-3.3-70B
+  # on one 96 GB card. Nothing else on the surface this service calls changed.
+  "interp-engine[vllm,quant]==1.8.0",
 ]
 
 [dependency-groups]

apps/inference/start.py

@@ -85,6 +85,20 @@ def parse_args():
         help="Model dtype: 'auto' loads each checkpoint's native dtype (gemma/qwen bf16, gpt2 fp32); "
         "pass e.g. float32/bfloat16 to force one.",
     )
+    parser.add_argument(
+        "--quantization",
+        default=None,
+        help="Quantize a wider checkpoint at load, with no calibration step: 'fp8' (vLLM backends), "
+        "'bnb-4bit' (vLLM or eager) or 'bnb-8bit' (eager). Embeddings stay at --model_dtype; the linear "
+        "layers narrow. A checkpoint that already ships quantized needs no flag. Needs interp-engine "
+        ">= 1.8. Env MODEL_QUANTIZATION overrides this flag.",
+    )
+    parser.add_argument(
+        "--kv_cache_dtype",
+        default=None,
+        help="Dtype of vLLM's paged KV cache: 'auto' (the model dtype) or 'fp8', which halves the cache "
+        "and about doubles the context that fits. vLLM backends only. Env KV_CACHE_DTYPE overrides this flag.",
+    )
     parser.add_argument(
         "--sae_dtype",
         default="bfloat16",
@@ -271,6 +285,10 @@ def main():
         os.environ["SAE_SETS"] = json.dumps([] if args.no_saes else args.sae_sets)
     if "MODEL_DTYPE" not in os.environ:
         os.environ["MODEL_DTYPE"] = args.model_dtype
+    if "MODEL_QUANTIZATION" not in os.environ and args.quantization is not None:
+        os.environ["MODEL_QUANTIZATION"] = args.quantization
+    if "KV_CACHE_DTYPE" not in os.environ and args.kv_cache_dtype is not None:
+        os.environ["KV_CACHE_DTYPE"] = args.kv_cache_dtype
     if "SAE_DTYPE" not in os.environ:
         os.environ["SAE_DTYPE"] = args.sae_dtype
     if "TOKEN_LIMIT" not in os.environ:

apps/inference/tests/unit/test_load_precision.py

@@ -0,0 +1,93 @@
+"""``--quantization`` and ``--kv_cache_dtype``: from start.py flag to the ``load_model`` call.
+
+Both are one name on ``load_model`` (interp-engine >= 1.8) and this server passes them through by
+that name, so the tests here are about the edges: a pod that sets neither makes the same call it
+always did, an engine too old to take them is refused with the pin to bump rather than dying in a
+constructor, and the KV cache dtype reaches the serving-limit arithmetic as a byte width.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+from neuronpedia_inference import server
+from neuronpedia_inference.server import _kv_cache_dtype_for_sizing, _load_precision_kwargs
+
+START_PY = Path(server.__file__).resolve().parents[1] / "start.py"
+
+
+def _new_load_model(hf_model_id, *, backend="auto", dtype="auto", quantization="", kv_cache_dtype="auto", **kw):
+    """The 1.8 signature, so these tests do not depend on which engine the venv holds."""
+
+
+def _old_load_model(hf_model_id, *, backend="auto", dtype="auto", **backend_kwargs):
+    """The 1.7 signature: the two names would fall into ``**backend_kwargs``."""
+
+
+def test_a_pod_that_sets_neither_adds_nothing_to_the_call(monkeypatch):
+    monkeypatch.setattr(server, "load_model", _old_load_model)
+    assert _load_precision_kwargs("", "auto") == {}
+    assert _load_precision_kwargs("", "") == {}
+
+
+def test_both_names_reach_load_model_as_load_model_spells_them(monkeypatch):
+    monkeypatch.setattr(server, "load_model", _new_load_model)
+    assert _load_precision_kwargs("fp8", "fp8") == {"quantization": "fp8", "kv_cache_dtype": "fp8"}
+    assert _load_precision_kwargs("bnb-4bit", "auto") == {"quantization": "bnb-4bit"}
+    assert _load_precision_kwargs("", "fp8") == {"kv_cache_dtype": "fp8"}
+
+
+def test_the_pinned_engine_takes_both_names():
+    """Against the engine the venv really holds, so a pin rolled back below 1.8 fails here first."""
+    assert _load_precision_kwargs("fp8", "fp8") == {"quantization": "fp8", "kv_cache_dtype": "fp8"}
+
+
+def test_an_engine_too_old_for_the_flag_is_refused_with_the_pin_to_bump(monkeypatch):
+    """On 1.7 the name would land in a constructor as an opaque TypeError, after the pod was paid for."""
+    monkeypatch.setattr(server, "load_model", _old_load_model)
+    with pytest.raises(ValueError, match="quantization needs interp-engine >= 1.8"):
+        _load_precision_kwargs("fp8", "auto")
+    with pytest.raises(ValueError, match="kv_cache_dtype needs interp-engine >= 1.8"):
+        _load_precision_kwargs("", "fp8")
+
+
+def test_the_kv_cache_is_priced_at_its_own_dtype_not_the_models():
+    """An fp8 cache holds twice the tokens; sizing it at bf16 would admit half the concurrency."""
+    assert _kv_cache_dtype_for_sizing("bfloat16", "auto") == "bfloat16"
+    assert _kv_cache_dtype_for_sizing("bfloat16", "") == "bfloat16"
+    assert _kv_cache_dtype_for_sizing("bfloat16", "fp8") == "float8"
+    assert _kv_cache_dtype_for_sizing("bfloat16", "fp8_e4m3") == "float8"
+
+
+def test_start_py_exports_the_flags_as_the_env_vars_args_py_reads():
+    """The whole start.py -> env -> args.py handoff, run the way a pod runs it."""
+    env = {k: v for k, v in os.environ.items() if k not in ("MODEL_QUANTIZATION", "KV_CACHE_DTYPE")}
+    code = (
+        "import os, sys; sys.argv = ['start.py', '--quantization', 'fp8', '--kv_cache_dtype', 'fp8', "
+        "'--list_models']; import start; "
+        "import neuronpedia_inference.args as a; a.list_available_options = lambda: None; "
+        "start.main(); print(os.environ['MODEL_QUANTIZATION'], os.environ['KV_CACHE_DTYPE'])"
+    )
+    out = subprocess.run(
+        [sys.executable, "-c", code], cwd=START_PY.parent, env=env, capture_output=True, text=True, check=True
+    )
+    assert out.stdout.split()[-2:] == ["fp8", "fp8"]
+
+
+def test_args_py_defaults_mean_as_stored_and_the_model_dtype(monkeypatch):
+    monkeypatch.delenv("MODEL_QUANTIZATION", raising=False)
+    monkeypatch.delenv("KV_CACHE_DTYPE", raising=False)
+    from neuronpedia_inference.args import parse_env_and_args
+
+    args = parse_env_and_args()
+    assert args.quantization == ""
+    assert args.kv_cache_dtype == "auto"
+    monkeypatch.setenv("MODEL_QUANTIZATION", "fp8")
+    monkeypatch.setenv("KV_CACHE_DTYPE", "fp8")
+    args = parse_env_and_args()
+    assert (args.quantization, args.kv_cache_dtype) == ("fp8", "fp8")