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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: interp-engine
-Version: 1.8.0
+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.8.0 → 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.8.0 → 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.8.0 → 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.8.0 → 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.8.0 → 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.8.0 → 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.8.0 → 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.8.0 → 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.8.0 → interp_engine-1.9.0}/interp_engine/load.py

@@ -152,9 +152,10 @@ def load_model(
             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

{interp_engine-1.8.0 → interp_engine-1.9.0}/interp_engine/memory.py

@@ -902,8 +902,9 @@ class WeightBytes:
     on_disk_bytes: int
     #: The dtype the checkpoint is stored in, e.g. ``bfloat16``. Empty when unknown.
     stored_dtype: str = ""
-    #: ``quantization_config.quant_method`` when there is one: ``mxfp4``, ``fp8``,
-    #: ``compressed-tensors``, ``awq``, ``gptq``. Empty on an unquantized checkpoint.
+    #: ``quantization_config.quant_method`` when there is one: ``mxfp4``, ``fp8``, ``awq``, ``gptq``.
+    #: A ``compressed-tensors`` container carries the width its ``config_groups`` declare, as
+    #: ``compressed-tensors int4`` or ``compressed-tensors fp8``. Empty on an unquantized checkpoint.
     quant_method: str = ""
     #: ``config.expert_dtype`` when an MoE checkpoint stores its routed experts narrower than the rest
     #: of itself, e.g. ``fp4`` beside a ``fp8`` scheme. Empty everywhere else.
@@ -1050,6 +1051,9 @@ def _quant_method(config: Any) -> tuple[str, str]:
         elif q is not None:
             method = str(getattr(q, "quant_method", "") or getattr(q, "quantization", "") or "")
         if method:
+            scheme = _compressed_tensors_scheme(q) if method.lower() == "compressed-tensors" else ""
+            if scheme:
+                method = f"{method} {scheme}"
             break
     stored = ""
     for holder in (config, text_config(config)):
@@ -1062,6 +1066,41 @@ def _quant_method(config: Any) -> tuple[str, str]:
     return method.lower(), stored
 
 
+def _compressed_tensors_scheme(quantization_config: Any) -> str:
+    """The weight format a ``compressed-tensors`` config really declares: ``int4``, ``int8``, ``fp8``.
+
+    ``compressed-tensors`` is a container, and the width it packs at sits one level down in
+    ``config_groups.*.weights`` as ``num_bits`` and ``type``. Read as a bare ``compressed-tensors``,
+    Kimi-K2.6's W4A16 experts are priced at one byte a parameter -- half the true dequantized size,
+    and the wrong answer to whether an eager load fits. The narrowest group wins, since on a mixed
+    checkpoint that is what the bulk of the weights are packed at. Empty when no group says.
+    """
+    groups = _get(quantization_config, "config_groups")
+    if isinstance(groups, dict):
+        groups = list(groups.values())
+    if not isinstance(groups, (list, tuple)):
+        return ""
+    best: tuple[int, str] | None = None
+    for group in groups:
+        weights = _get(group, "weights")
+        if weights is None:
+            continue
+        bits = _get(weights, "num_bits")
+        if not isinstance(bits, int) or bits not in (4, 8):
+            continue
+        kind = "fp" if str(_get(weights, "type") or "int").lower() == "float" else "int"
+        if best is None or bits < best[0]:
+            best = (bits, f"{kind}{bits}")
+    return best[1] if best else ""
+
+
+def _get(holder: Any, key: str) -> Any:
+    """One field of a config that may be a dict or an object."""
+    if isinstance(holder, dict):
+        return holder.get(key)
+    return getattr(holder, key, None)
+
+
 def _expert_dtype(config: Any) -> str:
     """The dtype of the ROUTED EXPERT weights, which need not be the checkpoint's dtype or its scheme.
 
@@ -1403,6 +1442,10 @@ class ModelMemoryFacts:
     v_head_dim: int
     vocab_size: int
     intermediate_size: int
+    #: What an MLA trunk caches per token per layer instead of K and V: the latent plus the RoPE'd
+    #: key part, 576 on DeepSeek-V3 and Kimi-K2. 0 where the cache holds K and V. One row, not one
+    #: per head, so it is also what says the cache cannot be sharded (:func:`kv_shards`).
+    kv_latent_width: int = 0
     #: Routed experts per sparse layer; 0 on a dense trunk. Here for the static tap widths --
     #: ``router_logits`` is as wide as the expert bank -- and because it is what says a layer's MLP is
     #: a fused kernel rather than three Linears, which decides whether ``mlp_act`` exists at all.
@@ -1431,7 +1474,11 @@ class ModelMemoryFacts:
     def kv_width(self) -> int:
         """KV-cache elements per token per layer, K and V together."""
         return kv_cache_width(
-            n_kv_heads=self.n_kv_heads, head_dim=self.head_dim, v_head_dim=self.v_head_dim, d_model=self.d_model
+            n_kv_heads=self.n_kv_heads,
+            head_dim=self.head_dim,
+            v_head_dim=self.v_head_dim,
+            d_model=self.d_model,
+            latent_width=self.kv_latent_width,
         )
 
     @property
@@ -1494,17 +1541,25 @@ class ModelMemoryFacts:
         return self.n_layers > 0 and self.d_model > 0
 
 
-def kv_cache_width(*, n_kv_heads: int = 0, head_dim: int = 0, v_head_dim: int = 0, d_model: int = 0) -> int:
+def kv_cache_width(
+    *, n_kv_heads: int = 0, head_dim: int = 0, v_head_dim: int = 0, d_model: int = 0, latent_width: int = 0
+) -> int:
     """KV-cache elements per token per layer, K and V together, for the whole model.
 
+    ``latent_width`` wins when set: an MLA trunk caches one latent row per token in place of K and V
+    (:func:`facts.kv_latent_width`), and its head dims describe tensors that never reach the cache.
+    Kimi-K2.6 at TP=8 on 8x H200 built 2.5x the tokens the head-dim figure predicted.
+
     ``2 * d_model`` -- the fallback when a caller has no head dims -- is the *pre-GQA* worst case and
     is wrong by 8x on the models where sizing is tight: Llama-3.3-70B caches 8 kv heads of 128 rather
-    than 8192, and a DeepSeek MLA trunk one 512-wide latent head rather than 4096. That factor is the
-    difference between fitting a default capture size and refusing to start, so pass the head dims.
+    than 8192. That factor is the difference between fitting a default capture size and refusing to
+    start, so pass the head dims.
 
     Mirrors ``vllm_capture/static.kv_cache_width`` and is checked against it in the tests; this copy
     exists so that sizing a model needs neither torch nor vLLM.
     """
+    if int(latent_width) > 0:
+        return int(latent_width)
     heads = max(int(n_kv_heads), 0)
     if heads <= 0 or int(head_dim) <= 0:
         return 2 * max(int(d_model), 1)
@@ -1517,8 +1572,8 @@ def kv_shards(facts: ModelMemoryFacts, num_gpus: int) -> int:
     **Not the rank count**, which is the assumption to get rid of. vLLM shards the cache by *KV
     head*, so how far it divides is a property of the attention shape rather than of the machine:
     Llama-3.3-70B's 8 KV heads go 2-per-rank at TP=4 and the cache really is a quarter on each card,
-    while a DeepSeek MLA trunk caches a single 512-wide latent head, which cannot be cut at all --
-    vLLM replicates it, and four cards hold four copies of the same cache.
+    while a DeepSeek MLA trunk caches a single 576-wide latent row (``kv_latent_width``), which
+    cannot be cut at all -- vLLM replicates it, and four cards hold four copies of the same cache.
 
     ``min`` covers both ends, including the case past the second one: vLLM pads a head count up to
     the rank count by duplicating heads, so 8 heads across 16 ranks still costs what 8 ranks cost
@@ -1530,7 +1585,7 @@ def kv_shards(facts: ModelMemoryFacts, num_gpus: int) -> int:
     lack.
     """
     tp = max(int(num_gpus), 1)
-    if facts.n_kv_heads <= 0 or facts.head_dim <= 0:
+    if facts.kv_latent_width > 0 or facts.n_kv_heads <= 0 or facts.head_dim <= 0:
         return 1
     return min(tp, facts.n_kv_heads)
 
@@ -2206,14 +2261,15 @@ def estimate(
     warnings: list[str] = []
     advice: list[str] = []
 
-    # Every row in `gpu-sizer/VERIFIED.md` was measured on one card. The single-GPU arithmetic is
-    # calibrated against hardware; how it divides across ranks is not, and the two terms it is most
-    # likely to be wrong about -- the weights, which TP shards unevenly, and the cache, which it shards
-    # for some attention shapes and replicates for others -- are the two largest.
+    # Tensor parallelism has two hardware measurements behind it (`gpu-sizer/VERIFIED.md`): Qwen3.8-27B
+    # at TP=2 on 2x A40 and Kimi-K2.6 at TP=8 on 8x H200, where the per-rank weight split and the margin
+    # outside the pool came out as priced. Every other count is still arithmetic. The Kimi run is also
+    # what showed the cache term had to price an MLA trunk's latent row rather than its heads.
     if tp > 1:
         warnings.append(
-            f"{tp}-GPU figures are unverified: every configuration measured so far ran on a single card, "
-            f"so tensor parallelism here is arithmetic no hardware has checked"
+            f"{tp}-GPU figures rest on two measurements: tensor parallelism has been checked on hardware "
+            f"at num_gpus=2 (Qwen3.8-27B on 2x A40) and num_gpus=8 (Kimi-K2.6 on 8x H200), where the "
+            f"weight split and the margin outside the pool held; other counts are still arithmetic"
         )
 
     # Only the eager backend expands a quantized checkpoint to the requested dtype; vLLM reads the same
@@ -2693,6 +2749,7 @@ def model_memory_facts(
         v_head_dim=f.v_head_dim,
         vocab_size=f.vocab_size,
         intermediate_size=intermediate,
+        kv_latent_width=f.kv_latent_width,
         n_experts=f.n_experts,
         layer_types=f.layer_types,
         sliding_window=f.sliding_window,