interp-engine 1.7.3__tar.gz → 1.9.0__tar.gz

{interp_engine-1.7.3 → interp_engine-1.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: interp-engine
-Version: 1.7.3
+Version: 1.9.0
 Summary: A fast, standardized, and easy to use interpretability engine.
 Project-URL: Homepage, https://github.com/decoderesearch/interp-engine
 Project-URL: Repository, https://github.com/decoderesearch/interp-engine

{interp_engine-1.7.3 → interp_engine-1.9.0}/benchmarks/probe.py

@@ -21,7 +21,11 @@ class EnvStamp:
     """What the numbers were produced on. Recorded per run, so a stale result is identifiable."""
 
     gpu_name: str = ""
+    """The card, prefixed with the count when the model was sharded across more than one
+    (``2x NVIDIA A40``), so every view keyed on this field says how many cards the number is from."""
     gpu_total_gib: float = 0.0
+    """Per card, not summed: the figure a reader compares against a card's spec sheet."""
+    gpu_count: int = 1
     driver_version: str = ""
     cuda_version: str = ""
     torch_version: str = ""
@@ -76,11 +80,13 @@ def _driver_version() -> str:
     return out.stdout.strip().splitlines()[0].strip() if out.stdout.strip() else "unknown"
 
 
-def env_stamp(index: int = 0) -> EnvStamp:
+def env_stamp(index: int = 0, num_gpus: int = 1) -> EnvStamp:
     props = torch.cuda.get_device_properties(index)
+    name = torch.cuda.get_device_name(index)
     return EnvStamp(
-        gpu_name=torch.cuda.get_device_name(index),
+        gpu_name=f"{num_gpus}x {name}" if num_gpus > 1 else name,
         gpu_total_gib=props.total_memory / GIB,
+        gpu_count=num_gpus,
         driver_version=_driver_version(),
         cuda_version=torch.version.cuda or "unknown",
         torch_version=torch.__version__,

{interp_engine-1.7.3 → interp_engine-1.9.0}/benchmarks/run_all.sh

@@ -11,6 +11,9 @@
 #   bash benchmarks/run_all.sh --models gemma-2-2b,qwen3-4b --variants eager,vllm
 #   bash benchmarks/run_all.sh --workloads generate,capture_mid --no-report
 #   bash benchmarks/run_all.sh --gpu-memory-utilization 0.7   # smaller card, or more worker scratch
+#   bash benchmarks/run_all.sh --num-gpus 1 --models qwen3.8-27b   # pin a multi-card box to one card
+#
+# Every cell is sharded across every visible CUDA card unless --num-gpus says otherwise.
 #   BENCH_PYTHON=/path/to/venv/bin/python bash benchmarks/run_all.sh
 #
 # With no --models, the sweep runs every model in the spec that this card can hold and names the ones
@@ -43,6 +46,7 @@ MODELS=""
 VARIANTS=""
 WORKLOADS=""
 GPU_MEM_UTIL=""
+NUM_GPUS=""
 RUN_REPORT=1
 SKIP_EXISTING=0
 
@@ -52,6 +56,7 @@ while [[ $# -gt 0 ]]; do
     --variants) VARIANTS="$2"; shift 2 ;;
     --workloads) WORKLOADS="$2"; shift 2 ;;
     --gpu-memory-utilization) GPU_MEM_UTIL="$2"; shift 2 ;;
+    --num-gpus) NUM_GPUS="$2"; shift 2 ;;
     --python) PYTHON="$2"; shift 2 ;;
     --no-report) RUN_REPORT=0; shift ;;
     # Resume a sweep that was interrupted. Off by default: a normal rerun should replace stale
@@ -70,13 +75,30 @@ if ! "$PYTHON" -c 'import interp_engine' 2>/dev/null; then
   exit 2
 fi
 
-# Total VRAM on the card this sweep will use, in GiB, or empty if there is no nvidia-smi to ask.
-# GiB rather than the vendor's GB, to match `min_gpu_gib` in the spec and the `gpu_total_gib` every
-# cell records -- a "180 GB" B200 reads as 179 GiB here, and comparing the two units is how a row
-# gets dropped on the one card that fits it.
+# No --num-gpus means every card the process can see: CUDA_VISIBLE_DEVICES when set, else what
+# nvidia-smi lists, else 1.
+if [[ -z "$NUM_GPUS" ]]; then
+  if [[ -n "${CUDA_VISIBLE_DEVICES+x}" ]]; then
+    NUM_GPUS=$(echo "$CUDA_VISIBLE_DEVICES" | tr ',' '\n' | grep -c .)
+  else
+    NUM_GPUS=$(nvidia-smi -L 2>/dev/null | grep -c '^GPU')
+  fi
+  (( NUM_GPUS > 0 )) || NUM_GPUS=1
+fi
+if ! [[ "$NUM_GPUS" =~ ^[1-9][0-9]*$ ]]; then
+  echo "error: --num-gpus must be a positive integer (got '$NUM_GPUS')" >&2
+  exit 2
+fi
+
+# Total VRAM this sweep will use, in GiB, or empty if there is no nvidia-smi to ask: the first
+# --num-gpus cards summed, since that is what a sharded load has to fit into. GiB rather than the
+# vendor's GB, to match `min_gpu_gib` in the spec and the `gpu_total_gib` every cell records -- a
+# "180 GB" B200 reads as 179 GiB here, and comparing the two units is how a row gets dropped on the
+# one card that fits it.
 gpu_total_gib() {
   local mib
-  mib="$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null | head -1 | tr -d ' ')"
+  mib="$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null \
+    | head -n "$NUM_GPUS" | tr -d ' ' | awk '{ s += $1 } END { if (NR) print s }')"
   [[ -z "$mib" ]] && return 0
   awk -v mib="$mib" 'BEGIN { printf "%.1f", mib / 1024 }'
 }
@@ -171,6 +193,7 @@ for pair in "${PAIRS[@]}"; do
   args=(--model "$model" --variant "$variant")
   [[ -n "$WORKLOADS" ]] && args+=(--workloads "$WORKLOADS")
   [[ -n "$GPU_MEM_UTIL" ]] && args+=(--gpu-memory-utilization "$GPU_MEM_UTIL")
+  (( NUM_GPUS > 1 )) && args+=(--num-gpus "$NUM_GPUS")
 
   # TOKENIZERS_PARALLELISM: the tokenizer is forked by vLLM's workers and warns on every cell
   # otherwise. VLLM_LOGGING_LEVEL: vLLM's per-step INFO logging would bury the workload lines.
@@ -200,6 +223,7 @@ if (( RUN_REPORT )); then
   [[ -n "$VARIANTS" ]] && sweep_cmd="$sweep_cmd --variants $VARIANTS"
   [[ -n "$WORKLOADS" ]] && sweep_cmd="$sweep_cmd --workloads $WORKLOADS"
   [[ -n "$GPU_MEM_UTIL" ]] && sweep_cmd="$sweep_cmd --gpu-memory-utilization $GPU_MEM_UTIL"
+  (( NUM_GPUS > 1 )) && sweep_cmd="$sweep_cmd --num-gpus $NUM_GPUS"
   "$PYTHON" -m benchmarks.report_bench --sweep-command "$sweep_cmd"
 fi
 

{interp_engine-1.7.3 → interp_engine-1.9.0}/benchmarks/run_bench.py

@@ -209,9 +209,13 @@ def _load_kwargs(
     m: ModelSpec,
     *,
     gpu_memory_utilization: float = GPU_MEMORY_UTILIZATION,
+    num_gpus: int = 1,
 ) -> dict[str, Any]:
     """Variant kwargs plus the per-backend knobs that only one backend accepts.
 
+    ``num_gpus`` is ``load_model``'s own knob and goes through as is: tensor parallelism on vLLM,
+    accelerate's layer placement on eager, where it takes the place of ``device="cuda"``.
+
     ``max_model_len`` and ``gpu_memory_utilization`` are vLLM-only, and passing them to the eager
     constructor would raise. Kept here rather than duplicated into every vLLM variant so the variant
     table stays about the thing being varied.
@@ -234,6 +238,8 @@ def _load_kwargs(
     kwargs = dict(v.kwargs)
     if kwargs.get("static_writes") == STEER_WRITES:
         kwargs["static_writes"] = [_steer_site(v, m)]
+    if num_gpus > 1:
+        kwargs["num_gpus"] = num_gpus
     # Every vLLM backend, not just the hooked one: these are engine settings, and a graph variant that
     # silently lost `max_model_len` would be measured against a different context than the row beside
     # it. `VLLM_BACKENDS` is the engine's own list, so a fourth backend is covered the day it lands.
@@ -270,9 +276,10 @@ def _load_kwargs(
             if isinstance(declared_model_kwargs, dict):
                 merged.update(declared_model_kwargs)
             kwargs["model_kwargs"] = merged
-        # A device_map places the weights itself, and `load_model` drops `device` when one is given.
-        # Setting it anyway would put a `device="cuda"` in the recorded kwargs that had no effect.
-        if "device_map" not in kwargs:
+        # A device_map places the weights itself, and `load_model` drops `device` when one is given
+        # -- as it does at `num_gpus > 1`, which is a device_map. Setting it anyway would put a
+        # `device="cuda"` in the recorded kwargs that had no effect.
+        if "device_map" not in kwargs and num_gpus == 1:
             kwargs.setdefault("device", "cuda")
     return kwargs
 
@@ -284,10 +291,11 @@ async def run_cell(
     *,
     command: str,
     gpu_memory_utilization: float = GPU_MEMORY_UTILIZATION,
+    num_gpus: int = 1,
 ) -> dict[str, Any]:
     from interp_engine import load_model
 
-    stamp = env_stamp()
+    stamp = env_stamp(num_gpus=num_gpus)
     record: dict[str, Any] = {
         "schema": SCHEMA,
         "started_at": datetime.now(UTC).isoformat(timespec="seconds"),
@@ -298,7 +306,7 @@ async def run_cell(
         "env": dataclasses.asdict(stamp),
         "workloads": {},
     }
-    kwargs = _load_kwargs(variant, model_spec, gpu_memory_utilization=gpu_memory_utilization)
+    kwargs = _load_kwargs(variant, model_spec, gpu_memory_utilization=gpu_memory_utilization, num_gpus=num_gpus)
     record["variant"]["kwargs"] = {k: str(val) for k, val in kwargs.items()}
     record["model"]["native_dtype"] = _native_dtype(model_spec.hf_id)
 
@@ -392,8 +400,18 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
             f"model's own declared fraction, or {GPU_MEMORY_UTILIZATION} where it has none"
         ),
     )
+    p.add_argument(
+        "--num-gpus",
+        type=int,
+        default=1,
+        help="shard the model across this many cards (vLLM tensor parallelism; accelerate placement "
+        "on eager). Stamped on the cell, since a two-card number is not a one-card number",
+    )
     p.add_argument("--list", action="store_true", help="print the models, variants and workloads and exit")
-    return p.parse_args(argv)
+    args = p.parse_args(argv)
+    if args.num_gpus < 1:
+        p.error("--num-gpus must be at least 1")
+    return args
 
 
 def main(argv: list[str] | None = None) -> int:
@@ -455,6 +473,7 @@ def main(argv: list[str] | None = None) -> int:
             workload_keys,
             command=command,
             gpu_memory_utilization=bench_spec.gpu_memory_utilization_for(model_spec.key, args.gpu_memory_utilization),
+            num_gpus=args.num_gpus,
         )
     )
 

{interp_engine-1.7.3 → interp_engine-1.9.0}/docs/AGENT_INTEGRATION.md

@@ -180,7 +180,7 @@ Canonical names, with the layer after a dot: `resid_post.10`. Extra coordinates
 | `mlp_act`, `mlp_pre`, `mlp_pre_linear` | MLP internals, `d_mlp` wide | **eager only** |
 | `router_logits` | MoE routing scores, every expert | both |
 | `expert_weights`, `expert_indices` | the top-k the router selected, and its weights | **eager only** |
-| the QK-norm points | inside the attention module | both, single-GPU only (head-sharded) |
+| the QK-norm points | inside the attention module | both (head-sharded; gathered across TP ranks) |
 
 Eager-only is not an omission: vLLM's fused MLP and MoE kernels compute those tensors inline, so
 there is no module boundary to hook. This table is the working subset;
@@ -191,7 +191,8 @@ Attention is the one row that reads "both" with a caveat. No boundary holds a sc
 backend, so `capture_attention(model, tokens, layers)` is how you ask, and it returns the same
 `{layer: {"scores", "probs", "value"}}` either way — from `output_attentions` on eager (which needs
 the model loaded with `attn_implementation="eager"`) and from an off-kernel recompute over captured
-post-RoPE q/k on vLLM (single-GPU only). Different code paths, same contract; `value` there is the
+post-RoPE q/k on vLLM (gathered across ranks under tensor parallelism, so it works at any
+`num_gpus`). Different code paths, same contract; `value` there is the
 per-head, family-scaled tensor satisfying `probs @ value == z`, not the raw projection output.
 
 [ENGINE_HOOK_MAPPINGS.md](ENGINE_HOOK_MAPPINGS.md) is the full dictionary across all three
@@ -228,9 +229,10 @@ an obvious one.
    always, on every configuration — the unembed happens in another process. Eager can do it, but only
    if the model was loaded with `requires_grad=True`. Gate on `model.grad_support`, not on backend
    name. See [GRADIENTS.md](GRADIENTS.md).
-5. **vLLM with `num_gpus > 1` serves no `z`, no DFA and no attention recompute.** Heads are sharded
-   across ranks, so rank 0 holds a slice; the engine refuses rather than returning it. Use one GPU, or
-   the eager backend, for per-head work.
+5. **vLLM with `num_gpus > 1` serves the same points as one GPU.** Heads and MLP neurons are
+   sharded across ranks, and the worker gathers `z`, `value`, `mlp_act`, the QK-norm points and the
+   q/k behind `capture_attention` back to full width at collect, so rank 0's payload is the whole
+   tensor. Every `collect_*` is a collective on every rank; do not call one from a single rank.
 6. **`await model.warmup()` before timing anything.** Construction is deliberately cheap and lazy on
    both backends — on vLLM nearly the whole load happens in `warmup()`, so without it your first
    request's latency is the load time.

{interp_engine-1.7.3 → interp_engine-1.9.0}/docs/INTERNALS.md

@@ -108,7 +108,15 @@ index names the layer HF's does, which nothing checked before and which would fa
 than raise. It needs `interp-engine[vllm]`, so it self-skips elsewhere; note that running it via
 `.venv-vllm/bin/python` needs that directory on `PATH` too, because vLLM shells out to `ninja` to
 build a sampler kernel at startup. `tests/test_vllm_wire_grammar.py` covers the same process
-boundary on CPU, over a synthetic demux.
+boundary on CPU, over a synthetic demux. `tests/test_multigpu.py` (`-m multigpu`, two CUDA cards)
+repeats the comparison at `num_gpus=2`: eager under accelerate's layer placement and vLLM under
+tensor parallelism, where `vllm_capture/_tp.py` gathers the head- and neuron-sharded points across
+ranks, against the one-card eager reference — captures, attention, steering, decode and the lens.
+One rule follows from tensor parallelism: a worker RPC refuses by *returning* a reason, never by
+raising. Under vLLM's multiprocess executor a raise inside `collective_rpc` is read from one rank
+while the others' replies stay queued, and the next RPC consumes those — so a single refused layer
+took every point of a TP=8 cell with it. `resolvable_points` and `resolvable_attn` are the ask-first
+calls; `tests/test_vllm_new_points.py` pins that they install nothing.
 
 **The docs are parsed, not maintained.** The point table in
 [SUPPORTED_POINTS.md](SUPPORTED_POINTS.md) and the footnote markers in

{interp_engine-1.7.3 → interp_engine-1.9.0}/docs/SUPPORTED_POINTS.md

@@ -14,12 +14,12 @@ nnsight and nnterp.
 | [`embeddings`][embeddings]                            | `d_model`               |  ✅   |  ✅  | trunk-level, so addressed with no layer index; distinct from `resid_pre` at layer 0 only where the trunk adds positional embeddings or scales the embedding                                                                               |
 | [`resid_pre`][resid_pre]                              | `d_model`               |  ✅   |  ✅  |                                                                                                                                                                                                                                           |
 | [`attn_in`][attn_in]                                  | `d_model`               |  ✅   |  ✅  |                                                                                                                                                                                                                                           |
-| [`q_norm_in`][q_norm_in] / [`q_norm_out`][q_norm_out] | `n_heads * head_dim`    |  ✅   |  ✅  | head-sharded, so single-GPU only                                                                                                                                                                                                          |
-| [`k_norm_in`][k_norm_in] / [`k_norm_out`][k_norm_out] | `n_kv_heads * head_dim` |  ✅   |  ✅  | head-sharded, so single-GPU only                                                                                                                                                                                                          |
-| [`value`][value]                                      | `n_heads * head_dim`    |  ✅   |  ✅  | head-sharded, so single-GPU only                                                                                                                                                                                                          |
+| [`q_norm_in`][q_norm_in] / [`q_norm_out`][q_norm_out] | `n_heads * head_dim`    |  ✅   |  ✅  | head-sharded; gathered across TP ranks at collect                                                                                                                                                                                         |
+| [`k_norm_in`][k_norm_in] / [`k_norm_out`][k_norm_out] | `n_kv_heads * head_dim` |  ✅   |  ✅  | head-sharded; gathered across TP ranks at collect                                                                                                                                                                                         |
+| [`value`][value]                                      | `n_heads * head_dim`    |  ✅   |  ✅  | head-sharded; gathered across TP ranks at collect                                                                                                                                                                                         |
 | [`attn_scores`][attn_scores]                          | `n_heads * query * key` |  ✅   |  ♻️  | no module boundary holds the pre-softmax matrix on **either** backend; vLLM rebuilds it from captured post-RoPE q/k                                                                                                                       |
 | [`attn_probs`][attn_probs]                            | `n_heads * query * key` |  ✅   |  ♻️  | fused paged attention never materializes the probabilities; same recompute                                                                                                                                                                |
-| [`z`][z]                                              | `n_heads * head_dim`    |  ✅   |  ✅  | head-sharded, so single-GPU only                                                                                                                                                                                                          |
+| [`z`][z]                                              | `n_heads * head_dim`    |  ✅   |  ✅  | head-sharded; gathered across TP ranks at collect                                                                                                                                                                                         |
 | [`attn_gate`][attn_gate]                              | `n_heads * head_dim`    |  ✅   |  ❌  | unimplemented — a real module on both trees                                                                                                                                                                                               |
 | [`attn_out`][attn_out]                                | `d_model`               |  ✅   |  ✅  |                                                                                                                                                                                                                                           |
 | [`attn_out_post`][attn_out_post]                      | `d_model`               |  ✅   |  ✅  |                                                                                                                                                                                                                                           |
@@ -27,7 +27,7 @@ nnsight and nnterp.
 | [`mlp_in`][mlp_in]                                    | `d_model`               |  ✅   |  ✅  |                                                                                                                                                                                                                                           |
 | [`mlp_pre`][mlp_pre]                                  | `d_mlp`                 |  ✅   |  ❌  | unreachable — vLLM fuses `gate_proj` and `up_proj` into one `gate_up_proj`, so neither branch is a module output                                                                                                                          |
 | [`mlp_pre_linear`][mlp_pre_linear]                    | `d_mlp`                 |  ✅   |  ❌  | as `mlp_pre`; gated MLPs only                                                                                                                                                                                                             |
-| [`mlp_act`][mlp_act]                                  | `d_mlp`                 |  ✅   |  ✅  | neuron-sharded, so single-GPU only                                                                                                                                                                                                        |
+| [`mlp_act`][mlp_act]                                  | `d_mlp`                 |  ✅   |  ✅  | neuron-sharded; gathered across TP ranks at collect                                                                                                                                                                                       |
 | [`router_logits`][router_logits]                      | `n_experts`             |  ✅   |  ✅  | replicated gate, so it survives tensor parallelism                                                                                                                                                                                        |
 | [`expert_weights`][expert_weights]                    | `n_experts`             |  ✅   |  ❌  | unreachable — the top-k happens inside the FusedMoE kernel, which returns the combined output with the selection never materialized                                                                                                       |
 | [`expert_indices`][expert_indices]                    | `n_experts`             |  ✅   |  ❌  | as `expert_weights`                                                                                                                                                                                                                       |
@@ -54,13 +54,16 @@ a fused kernel ate the tensor and no module boundary holds it. Ask the code rath
 you are branching on it — `points.vllm_hookable()` is the served set, `points.reason(name)` is the
 sentence for one refusal, and `model.points()` is what a loaded model has.
 
-## Tensor parallelism narrows the vLLM column further
+## Tensor parallelism does not narrow the vLLM column
 
-The capture path reads rank 0's payload alone, so a point whose last axis vLLM shards comes back as a
-slice: `z`, `value`, `mlp_act` and the four QK-norm points are refused on a multi-GPU pod rather than
-returned short, and so is the attention recompute (q/k/v are head-sharded). Everything `d_model` wide
-is all-reduced before the hook sees it, and `router_logits` comes off a replicated gate, so those are
-unaffected.
+The capture path reads rank 0's payload alone, and at `num_gpus > 1` vLLM shards `z`, `value`,
+`mlp_act`, the four QK-norm points and the q/k/v behind the attention recompute by head or by neuron.
+The worker gathers those across ranks at collect time (`interp_engine.vllm_capture._tp`), so rank 0
+hands back the same full-width tensor a single card would have, in the same layout; a KV head that
+vLLM replicates rather than shards is kept once. Everything `d_model` wide is all-reduced before the
+hook sees it, and `router_logits` comes off a replicated gate, so those never needed gathering.
+Verified against the eager reference at TP=2 by the validator (`validator/`, `NUM_GPUS=2`) and by
+`tests/test_multigpu.py`.
 
 ## The last seven rows need a hyper-connection trunk
 

{interp_engine-1.7.3 → interp_engine-1.9.0}/docs/USAGE.md

@@ -94,6 +94,16 @@ eager = load_model("google/gemma-2-2b-it", backend="eager", device="cuda", dtype
 served = load_model("meta-llama/Llama-3.1-8B", backend="vllm", gpu_memory_utilization=0.85)
 ```
 
+Two knobs narrow a bf16 checkpoint on load, with no calibration and no other repo:
+`quantization=` applies `"fp8"` (vLLM backends), `"bnb-4bit"` (either backend) or `"bnb-8bit"`
+(eager) to every linear layer, and `kv_cache_dtype="fp8"` halves vLLM's paged cache. Both take the
+same name on both backends and are refused, naming the alternative, where a backend cannot apply
+them. The [GPU sizer](../gpu-sizer/INPUTS.md) prices both and prints the argument it priced.
+
+```python
+narrow = load_model("meta-llama/Llama-3.3-70B-Instruct", backend="vllm", quantization="fp8", kv_cache_dtype="fp8")
+```
+
 **Construction is cheap and lazy on both backends; `warmup()` is where the cost lands.** On vLLM
 almost the entire load happens there, so call it before you time anything or a first request pays
 for the engine.

{interp_engine-1.7.3 → interp_engine-1.9.0}/interp_engine/capture.py

@@ -480,8 +480,8 @@ def capture_attention(
     pass rather than one rebuilt from the other.
 
     The eager arm needs the model loaded with eager attention for ``attn_probs``; the vLLM arm
-    recomputes off-kernel from captured post-RoPE q/k and is single-GPU only. Both refusals name
-    themselves.
+    recomputes off-kernel from captured post-RoPE q/k, gathered across ranks under tensor
+    parallelism. Both refusals name themselves.
     """
     if not isinstance(model, EagerModel):
         ids = as_token_ids(tokens, model=model, what="capture_attention")

{interp_engine-1.7.3 → interp_engine-1.9.0}/interp_engine/facts.py

@@ -1630,6 +1630,21 @@ def value_head_dim(cfg: Any, head_dim: int) -> int:
     return _first_int(cfg, ("v_head_dim",)) or head_dim
 
 
+def kv_latent_width(cfg: Any) -> int:
+    """Elements one token of KV cache holds per layer under MLA, or 0 where keys and values are cached.
+
+    DeepSeek's multi-head latent attention never caches K and V: it caches the ``kv_lora_rank``-wide
+    latent both are expanded from, plus the ``qk_rope_head_dim`` positional key part that rides beside
+    it -- one 576-wide row per token on DeepSeek-V3 and Kimi-K2, in place of 64 heads of K and V. A
+    cache sized from ``n_kv_heads * (head_dim + v_head_dim)`` is 21x too large there, and it is the
+    figure vLLM builds its pool from, so this is what a sizer has to charge. 0 on every non-MLA config.
+    """
+    rank = _first_int(cfg, ("kv_lora_rank",))
+    if not rank:
+        return 0
+    return rank + (_first_int(cfg, ("qk_rope_head_dim",)) or 0)
+
+
 def first_kv_shared_layer(cfg: Any, n_layers: int) -> int | None:
     """The first layer that reuses an earlier layer's keys/values, or None if none do."""
     shared = _first_int(cfg, ("num_kv_shared_layers",))
@@ -2020,6 +2035,9 @@ class ModelFacts:
     # The width of one value head, which differs from ``head_dim`` on MiMo-V2 and the DeepSeek MLA
     # families. See :func:`value_head_dim`; ``value`` and ``z`` are this wide per head, not ``head_dim``.
     v_head_dim: int = 0
+    # Elements per token per layer that an MLA trunk caches in place of K and V (the latent plus the
+    # RoPE'd key part); 0 where the cache holds K and V. See :func:`kv_latent_width`.
+    kv_latent_width: int = 0
     # First layer that reuses an earlier layer's keys/values and so has no k/v projection of its own
     # (Gemma-4); None when every layer computes its own.
     first_kv_shared_layer: int | None = None
@@ -2476,6 +2494,7 @@ def resolve_facts(config: Any, *, n_layers_fallback: int | None = None) -> Model
         global_kv_heads=_first_int(cfg, ("num_global_key_value_heads",)) or None,
         k_eq_v=bool(config_attr(cfg, "attention_k_eq_v", False)),
         v_head_dim=value_head_dim(cfg, head_dim),
+        kv_latent_width=kv_latent_width(cfg),
         first_kv_shared_layer=first_kv_shared_layer(cfg, n_layers),
         n_experts=n_experts(cfg),
         experts_per_token=_first_int(cfg, _EXPERTS_PER_TOKEN_FIELDS) or 0,

{interp_engine-1.7.3 → interp_engine-1.9.0}/interp_engine/load.py

@@ -50,12 +50,62 @@ def _declares_nothing(value: Any) -> bool:
     return not list(value)
 
 
+def _apply_load_precision(backend: str, quantization: str, kv_cache_dtype: str, backend_kwargs: dict[str, Any]) -> None:
+    """Turn ``quantization`` and ``kv_cache_dtype`` into what ``backend``'s constructor takes.
+
+    Both are one name on ``load_model`` and two different things underneath: vLLM quantizes through
+    an engine argument, transformers through a ``BitsAndBytesConfig``. The table in
+    :data:`interp_engine.memory.QUANTIZATIONS` says which backend applies which scheme, and a scheme
+    the backend cannot apply is refused here with that table's reason -- never passed on to become an
+    opaque ``TypeError`` from a constructor, and never dropped to load the checkpoint as stored.
+    """
+    from interp_engine.memory import QUANTIZATIONS, quantization_refusal
+
+    use_vllm = backend in VLLM_BACKENDS
+    if kv_cache_dtype not in ("auto", "", None):
+        if not use_vllm:
+            raise ValueError(
+                f"kv_cache_dtype={kv_cache_dtype!r} names the dtype of vLLM's paged KV cache, and "
+                f"backend={backend!r} has no such cache. Drop it, or use a vLLM backend."
+            )
+        extra = dict(backend_kwargs.get("extra_vllm_kwargs") or {})
+        extra.setdefault("kv_cache_dtype", kv_cache_dtype)
+        backend_kwargs["extra_vllm_kwargs"] = extra
+
+    if not quantization:
+        return
+    refused = quantization_refusal(quantization, backend)
+    if refused:
+        raise ValueError(f"quantization={quantization!r} on backend={backend!r}: {refused}")
+    scheme = QUANTIZATIONS[quantization]
+    if use_vllm:
+        extra = dict(backend_kwargs.get("extra_vllm_kwargs") or {})
+        if extra.get("quantization") not in (None, scheme.vllm_name):
+            raise ValueError(
+                f"quantization={quantization!r} asks vLLM for {scheme.vllm_name!r}, but extra_vllm_kwargs "
+                f"already names {extra['quantization']!r}. Pass one or the other."
+            )
+        extra["quantization"] = scheme.vllm_name
+        backend_kwargs["extra_vllm_kwargs"] = extra
+        return
+    if backend_kwargs.get("quantization_config") is not None:
+        raise ValueError(
+            f"quantization={quantization!r} builds a BitsAndBytesConfig, and quantization_config= was "
+            f"passed as well. Pass one or the other."
+        )
+    from transformers import BitsAndBytesConfig
+
+    backend_kwargs["quantization_config"] = BitsAndBytesConfig(**scheme.eager_config)
+
+
 def load_model(
     hf_model_id: str,
     *,
     backend: str = "auto",
     device: str | None = None,
     dtype: str = "auto",
+    quantization: str = "",
+    kv_cache_dtype: str = "auto",
     num_gpus: int = 1,
     trust_remote_code: bool | None = None,
     static_points: Any = None,
@@ -86,11 +136,26 @@ def load_model(
         device: Explicit device for the eager backend. None means let the ladder choose.
             Ignored by vLLM, which always initializes on CUDA.
         dtype: ``"auto"`` (the checkpoint's native precision) or an explicit
-            ``"float32"``/``"float16"``/``"bfloat16"``.
+            ``"float32"``/``"float16"``/``"bfloat16"``. This is the width the **activations** run
+            at, and the width an unquantized checkpoint's weights are held at. It is not how to ask
+            for a narrower checkpoint: vLLM rejects ``dtype="fp8"``, and transformers would store
+            fp8 weights with no kernel behind them. That is ``quantization``.
+        quantization: An on-load scheme from :data:`interp_engine.memory.QUANTIZATIONS`, applied to
+            a wider checkpoint as it loads, with no calibration step: ``"fp8"`` (vLLM only),
+            ``"bnb-4bit"`` (either backend) or ``"bnb-8bit"`` (eager only). Empty, the default,
+            loads the checkpoint as stored -- which for a repo that already ships quantized is the
+            right answer, since a quantizer cannot narrow what is already narrower. A scheme the
+            chosen backend cannot apply is refused, naming the one to use instead. Other vLLM
+            schemes still reach the engine through ``extra_vllm_kwargs={"quantization": ...}``.
+        kv_cache_dtype: vLLM's KV cache dtype -- ``"auto"`` (the model dtype, or the scheme the
+            checkpoint declares for its cache) or ``"fp8"``, which halves the cache and so roughly
+            doubles the context or concurrency a card holds. Refused on the eager backend, which
+            has no paged cache to set the dtype of.
         num_gpus: Shard across this many GPUs on one node -- vLLM ``tensor_parallel_size``,
-            eager accelerate ``device_map="auto"``. Note that vLLM with ``num_gpus > 1``
-            cannot serve per-head ``z`` or DFA, because attention heads are sharded across
-            ranks and the off-kernel recompute would only see one shard.
+            eager accelerate ``device_map="auto"``. The vLLM worker gathers the head- and
+            neuron-sharded points (``z``, ``value``, ``mlp_act``, the QK-norm points, the q/k
+            behind the attention recompute) across ranks at collect, so the served point set
+            and every tensor's width are the same as on one GPU.
         trust_remote_code: Passed to both the config probe and the backend. The default ``None``
             means "only where the checkpoint has no alternative": eager prefers a native
             transformers class over a checkpoint's bundled copy of one when both exist, since the
@@ -118,7 +183,8 @@ def load_model(
         ValueError: ``backend`` is not one of :data:`BACKENDS`; or ``static_points`` /
             ``static_writes`` was passed on a backend other than ``"vllm-static"``; or
             ``backend="vllm-static"`` declared no taps at all; or ``enforce_eager=True`` was
-            passed alongside a graph-replaying backend.
+            passed alongside a graph-replaying backend; or ``quantization`` / ``kv_cache_dtype``
+            asks the chosen backend for something it cannot apply.
         RuntimeError: a vLLM backend was requested but vLLM is not installed.
         GradientsUnsupported: ``requires_grad=True`` on a vLLM backend, which cannot
             provide gradients through its forward on any configuration.
@@ -182,6 +248,8 @@ def load_model(
         # installs nothing in Worker.load_model, and leaves hooks_available False.
         static_points, static_writes = [], None
 
+    _apply_load_precision(resolved, quantization, kv_cache_dtype, backend_kwargs)
+
     if use_vllm:
         require_vllm(f"backend={resolved!r} requested for {hf_model_id}")
         # `requires_grad` is an eager-only constructor kwarg, so on vLLM it would otherwise land as