actlens 0.1.0__py3-none-any.whl

actlens/__init__.py

@@ -0,0 +1,3 @@
+"""ActLens: interactive activation viewer backend."""
+
+__version__ = "0.1.0"

actlens/app.py

@@ -0,0 +1,318 @@
+"""FastAPI app: model management, runs, and windowed activation slices."""
+from __future__ import annotations
+
+import asyncio
+import gc
+import hmac
+import json
+import os
+import threading
+import time
+import uuid
+from contextlib import asynccontextmanager
+from collections import OrderedDict
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from typing import Callable
+
+from fastapi import FastAPI, HTTPException, Query, Request
+from fastapi.responses import FileResponse, PlainTextResponse, RedirectResponse
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel
+
+from . import slicing
+from .cache import ActivationCache
+from .capture import ATTN_ACT, Capture, Run
+from .providers import ActivationProvider, NNsightProvider
+from .wire import frame_response
+
+DEFAULT_MODEL = "Qwen/Qwen3-0.6B"
+PRESET_MODELS = [
+    {"id": "Qwen/Qwen3-0.6B", "label": "Qwen3-0.6B"},
+    {"id": "Qwen/Qwen2.5-0.5B", "label": "Qwen2.5-0.5B"},
+    {"id": "HuggingFaceTB/SmolLM2-360M", "label": "SmolLM2-360M"},
+    {"id": "openai-community/gpt2", "label": "GPT-2"},
+]
+MAX_RUNS = 3
+TOKEN_COOKIE = "actlens_token"
+DEFAULT_CACHE_MB = 2048
+CORPUS = json.loads((Path(__file__).parent / "corpus.json").read_text())
+_PACKAGED_DIST = Path(__file__).parent / "static"  # bundled into the wheel
+_DEV_DIST = Path(__file__).resolve().parents[2] / "frontend" / "dist"  # source checkout
+FRONTEND_DIST = _PACKAGED_DIST if (_PACKAGED_DIST / "index.html").is_file() else _DEV_DIST
+
+
+def cache_budget_bytes() -> int:
+    try:
+        mb = float(os.environ.get("ACTLENS_CACHE_MB", DEFAULT_CACHE_MB))
+    except ValueError:
+        mb = DEFAULT_CACHE_MB
+    return int(mb * 2**20)
+
+
+class ModelManager:
+    """Owns the single loaded model. All model work runs on one dedicated thread."""
+
+    def __init__(self, factory: Callable[[str, str, str], ActivationProvider] = NNsightProvider,
+                 cache_bytes: int | None = None):
+        self.factory = factory
+        self.executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="model")
+        self.provider: ActivationProvider | None = None
+        self.state = "idle"  # idle | loading | ready | error
+        self.target: str | None = None
+        self.error: str | None = None
+        self.runs: OrderedDict[str, Run] = OrderedDict()
+        self.cache = ActivationCache(cache_budget_bytes() if cache_bytes is None else cache_bytes)
+        self._lock = threading.Lock()
+
+    def status(self) -> dict:
+        return {
+            "state": self.state,
+            "model_id": self.provider.model_id if self.provider else self.target,
+            "target": self.target,
+            "error": self.error,
+            "info": self.provider.info if self.state == "ready" and self.provider else None,
+        }
+
+    def load(self, model_id: str, device: str = "auto", dtype: str = "float32") -> None:
+        with self._lock:
+            if self.state == "loading":
+                raise HTTPException(409, "A model is already loading")
+            self.state, self.target, self.error = "loading", model_id, None
+        self.executor.submit(self._load, model_id, device, dtype)
+
+    def _load(self, model_id: str, device: str, dtype: str) -> None:
+        try:
+            if self.provider is not None:
+                self.provider.close()
+                self.provider = None
+            self.runs.clear()
+            self.cache.clear()
+            gc.collect()
+            self.provider = self.factory(model_id, device, dtype)
+            self.state = "ready"
+        except Exception as e:  # surfaced to the UI via /api/status
+            self.error = f"{type(e).__name__}: {e}"
+            self.state = "error"
+
+    def _register(self, provider: ActivationProvider, text: str, max_tokens: int) -> Run:
+        t0 = time.perf_counter()
+        ids, tokens, truncated = provider.tokenize(text, max_tokens)
+        specs = {s.id: s for s in provider.activations()}
+        keys = ("n_layers", "hidden_size", "intermediate_size", "n_heads", "n_kv_heads", "head_dim")
+        return Run(run_id=uuid.uuid4().hex[:12], model_id=provider.model_id, text=text, token_ids=list(ids),
+                   tokens=list(tokens), truncated=truncated, specs=specs,
+                   model={k: provider.info.get(k) for k in keys}, elapsed_ms=(time.perf_counter() - t0) * 1000)
+
+    async def run(self, text: str, max_tokens: int) -> Run:
+        """Tokenize and register a run. No activations are captured yet."""
+        if self.state != "ready" or self.provider is None:
+            raise HTTPException(409, f"Model not ready (state: {self.state})")
+        loop = asyncio.get_running_loop()
+        try:
+            run = await loop.run_in_executor(self.executor, self._register, self.provider, text, max_tokens)
+        except ValueError as e:
+            raise HTTPException(400, str(e))
+        self.runs[run.run_id] = run
+        while len(self.runs) > MAX_RUNS:
+            old, _ = self.runs.popitem(last=False)
+            self.cache.drop_run(old)
+        return run
+
+    def get(self, run_id: str) -> Run:
+        run = self.runs.get(run_id)
+        if run is None:
+            raise HTTPException(404, "Unknown run (it may have been evicted); run the prompt again")
+        self.runs.move_to_end(run_id)
+        return run
+
+    def capture(self, run: Run, act: str) -> Capture:
+        """Blocking (call from a worker thread): return the cached capture, building it on the model thread if needed.
+        Concurrent callers for the same (run, act) share one forward pass."""
+        provider = self.provider
+        if provider is None or run.run_id not in self.runs or provider.model_id != run.model_id:
+            raise HTTPException(404, "Unknown run (it may have been evicted); run the prompt again")
+
+        def build() -> Capture:
+            arr = self.executor.submit(provider.capture, run.token_ids, act).result()
+            return Capture(act=act, arr=arr)
+
+        try:
+            cap = self.cache.get((run.run_id, act), build)
+        except HTTPException:
+            raise
+        except ValueError as e:
+            raise HTTPException(400, str(e))
+        except Exception as e:
+            raise HTTPException(500, f"Capture of {act!r} failed: {type(e).__name__}: {e}")
+        if run.run_id not in self.runs:  # run was evicted while we were capturing
+            self.cache.drop_run(run.run_id)
+        return cap
+
+
+class LoadRequest(BaseModel):
+    model_id: str
+    device: str = "auto"
+    dtype: str = "float32"
+
+
+class RunRequest(BaseModel):
+    text: str
+    max_tokens: int = 512
+
+
+def create_app(manager: ModelManager | None = None, autoload: bool = True, model_id: str = DEFAULT_MODEL,
+               device: str = "auto", dtype: str = "float32", token: str | None = None) -> FastAPI:
+    mgr = manager or ModelManager()
+
+    @asynccontextmanager
+    async def lifespan(_: FastAPI):
+        if autoload and mgr.state == "idle":
+            mgr.load(model_id, device, dtype)
+        yield
+
+    app = FastAPI(title="ActLens", lifespan=lifespan)
+    app.state.manager = mgr
+
+    if token:
+        # Access token for a server that is reachable from outside (e.g. a tunnel from Colab). Open
+        # `/?token=...` once: the token is stored in an HttpOnly cookie and the URL is cleaned up.
+        # Scripts can send `Authorization: Bearer <token>` instead.
+        def valid(candidate: str | None) -> bool:
+            return bool(candidate) and hmac.compare_digest(candidate.encode(), token.encode())
+
+        @app.middleware("http")
+        async def require_token(request: Request, call_next):
+            bearer = request.headers.get("authorization", "")
+            if valid(request.cookies.get(TOKEN_COOKIE)) or valid(bearer.removeprefix("Bearer ").strip()):
+                return await call_next(request)
+            if request.method == "GET" and valid(request.query_params.get("token")):
+                url = request.url.remove_query_params("token")
+                resp = RedirectResponse(url.path + (f"?{url.query}" if url.query else ""), status_code=303)
+                # Lax, not Strict: the link is usually opened from another site (Colab), and a Strict cookie is
+                # not sent on the redirect that follows a cross-site navigation.
+                https = "https" in (request.url.scheme, request.headers.get("x-forwarded-proto", ""))
+                resp.set_cookie(TOKEN_COOKIE, token, httponly=True, samesite="lax", secure=https)
+                return resp
+            return PlainTextResponse("ActLens: missing or invalid access token. Open the full link printed "
+                                     "by the server (it ends in ?token=...).", status_code=401)
+
+    def guarded(fn, *a, **kw):
+        try:
+            return fn(*a, **kw)
+        except ValueError as e:
+            raise HTTPException(400, str(e))
+
+    def token_capture(run_id: str, act: str, layer: int | None = None) -> Capture:
+        """Validate `act` (a token-kind activation of this run's model) and capture it lazily."""
+        run = mgr.get(run_id)
+        spec = run.specs.get(act)
+        if spec is None:
+            raise HTTPException(400, f"unknown act {act!r}; expected one of {list(run.specs)}")
+        if spec.kind != "token":
+            raise HTTPException(400, f"act {act!r} is an attention pattern; use the /attn endpoints")
+        if layer is not None and not 0 <= layer < spec.n_layers:
+            raise HTTPException(400, f"layer {layer} out of range for {act} (0..{spec.n_layers - 1})")
+        return mgr.capture(run, act)
+
+    def attn_capture(run_id: str) -> Capture:
+        run = mgr.get(run_id)
+        if ATTN_ACT not in run.specs:
+            raise HTTPException(400, "this model does not expose attention patterns")
+        return mgr.capture(run, ATTN_ACT)
+
+    # ----- models / corpus -----
+    @app.get("/api/status")
+    def status():
+        return {**mgr.status(), "presets": PRESET_MODELS}
+
+    @app.post("/api/models/load", status_code=202)
+    def load_model(req: LoadRequest):
+        mgr.load(req.model_id, req.device, req.dtype)
+        return mgr.status()
+
+    @app.get("/api/corpus")
+    def corpus():
+        return CORPUS
+
+    # ----- runs -----
+    @app.post("/api/run")
+    async def run(req: RunRequest):
+        r = await mgr.run(req.text, max(1, min(req.max_tokens, 1024)))
+        return {
+            "run_id": r.run_id, "model_id": r.model_id, "tokens": r.tokens, "token_ids": r.token_ids,
+            "truncated": r.truncated, "elapsed_ms": round(r.elapsed_ms, 1), "model": r.model,
+            "activations": [sp.info() for sp in r.specs.values()],
+        }
+
+    @app.get("/api/run/{run_id}/overview")
+    def overview(run_id: str, act: str = "resid_post", stat: str = "norm", dim: int = 0):
+        cap = token_capture(run_id, act)
+        arr, meta = guarded(slicing.overview, cap, stat, dim)
+        return frame_response(meta, arr)
+
+    @app.get("/api/run/{run_id}/slice")
+    def slice_(run_id: str, act: str, layer: int, t0: int = 0, t1: int = 64, d0: int = 0, d1: int = 128,
+               max_h: int = Query(512, ge=1, le=2048), max_w: int = Query(1024, ge=1, le=4096),
+               agg: str = "absmax", order: str = "natural"):
+        cap = token_capture(run_id, act, layer)
+        arr, meta = guarded(slicing.token_dim_slice, cap, layer, t0, t1, d0, d1, max_h, max_w, agg, order)
+        return frame_response(meta, arr)
+
+    @app.get("/api/run/{run_id}/profile")
+    def profile(run_id: str, act: str, layer: int, order: str = "natural", bins: int = Query(512, ge=8, le=4096)):
+        cap = token_capture(run_id, act, layer)
+        arr, meta = guarded(slicing.dim_profile, cap, layer, order, bins)
+        return frame_response(meta, arr)
+
+    @app.get("/api/run/{run_id}/stats")
+    def stats(run_id: str, act: str, layer: int, t0: int = 0, t1: int = 64, d0: int = 0, d1: int = 128,
+              order: str = "natural", clip: bool = False):
+        cap = token_capture(run_id, act, layer)
+        return guarded(slicing.token_region_stats, cap, layer, t0, t1, d0, d1, order, clip)
+
+    @app.get("/api/run/{run_id}/axis_stats")
+    def axis_stats(run_id: str, act: str, layer: int, axis: str = "channel", stat: str = "norm",
+                   t0: int = 0, t1: int = 64, d0: int = 0, d1: int = 128, order: str = "natural",
+                   clip: bool = False, top: int = 10, bins: int = 64):
+        cap = token_capture(run_id, act, layer)
+        return guarded(slicing.axis_stats, cap, layer, axis, stat, t0, t1, d0, d1, order, clip, top, bins)
+
+    @app.get("/api/run/{run_id}/attn")
+    def attn(run_id: str, layer: int, head: int = -1, q0: int = 0, q1: int = 100000, k0: int = 0, k1: int = 100000,
+             max_q: int = Query(512, ge=1, le=1024), max_k: int = Query(512, ge=1, le=1024), agg: str = "max"):
+        cap = attn_capture(run_id)
+        if not 0 <= layer < cap.n_layers:
+            raise HTTPException(400, f"layer {layer} out of range (0..{cap.n_layers - 1})")
+        arr, meta = guarded(slicing.attn_slice, cap, layer, head, q0, q1, k0, k1, max_q, max_k, agg)
+        return frame_response(meta, arr)
+
+    @app.get("/api/run/{run_id}/attn_overview")
+    def attn_overview(run_id: str, stat: str = "entropy"):
+        cap = attn_capture(run_id)
+        arr, meta = guarded(slicing.attn_overview, cap, stat)
+        return frame_response(meta, arr)
+
+    @app.get("/api/run/{run_id}/attn_stats")
+    def attn_stats(run_id: str, layer: int, head: int, q0: int = 0, q1: int = 100000, k0: int = 0, k1: int = 100000,
+                   clip: bool = False):
+        cap = attn_capture(run_id)
+        if not (0 <= layer < cap.n_layers and 0 <= head < cap.n_heads):
+            raise HTTPException(400, "layer/head out of range")
+        return guarded(slicing.attn_region_stats, cap, layer, head, q0, q1, k0, k1, clip)
+
+    # ----- static frontend (production build) -----
+    if (FRONTEND_DIST / "index.html").is_file():
+        app.mount("/assets", StaticFiles(directory=FRONTEND_DIST / "assets"), name="assets")
+
+        @app.get("/{path:path}")
+        def spa(path: str):
+            if path.startswith("api/"):
+                raise HTTPException(404)
+            f = FRONTEND_DIST / path
+            return FileResponse(f if path and f.is_file() else FRONTEND_DIST / "index.html")
+
+    return app
+
+
+app = create_app()

actlens/archs/__init__.py

@@ -0,0 +1,7 @@
+"""Architecture adapters: how each model family exposes its activations. See `base.ArchAdapter`."""
+from . import gpt2, llama  # noqa: F401  (importing registers the built-in adapters)
+from .base import ActDef, ArchAdapter, Dims, PreNormBlockAdapter, TraceCtx, get, has
+from .registry import ADAPTERS, load_plugins, register, resolve_adapter, supported_model_types
+
+__all__ = ["ActDef", "ArchAdapter", "Dims", "PreNormBlockAdapter", "TraceCtx", "ADAPTERS", "get", "has",
+           "load_plugins", "register", "resolve_adapter", "supported_model_types"]

actlens/archs/base.py

@@ -0,0 +1,153 @@
+"""The architecture-adapter contract.
+
+An adapter answers two questions about one family of Hugging Face models: *which* activations does it expose, and
+*how* is each one read from a block during a forward pass. Everything else (tokenizing, tracing, caching, the HTTP
+API, the UI) is architecture-neutral and only sees the resulting `ActivationSpec`s and arrays.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Callable
+
+import torch
+
+
+@dataclass(frozen=True)
+class Dims:
+    """Model dimensions the provider and the specs need. `inter` is None for models without an MLP width."""
+    n_layers: int
+    hidden: int
+    inter: int | None
+    n_heads: int
+    n_kv_heads: int
+    head_dim: int
+
+
+@dataclass
+class TraceCtx:
+    """Per-capture state handed to readers. `aux` holds values an adapter computes once before the blocks run
+    (see `ArchAdapter.prepare`), e.g. the RoPE cos/sin tables."""
+    n_tokens: int
+    aux: dict[str, Any] = field(default_factory=dict)
+
+    def host(self, t: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor:
+        return t.detach().to(dtype).cpu()
+
+    def tok(self, t: torch.Tensor) -> torch.Tensor:
+        """[1, T, C] or [1, T, heads, Dh] -> host float32 [T, C]."""
+        return self.host(t).reshape(self.n_tokens, -1)
+
+
+# (block, ctx) -> host tensor: [T, C] float32 for token activations, [H, T, T] float16 for attention patterns.
+# Runs inside an nnsight trace, so `block.<module>.input/output` are lazy proxies; readers must touch modules in
+# forward-execution order and use `ctx` to move data to the host.
+Reader = Callable[[Any, TraceCtx], torch.Tensor]
+
+
+@dataclass(frozen=True)
+class ActDef:
+    """How one activation is exposed on one architecture. `label`/`description` override the registry defaults
+    (use them to name the model's real modules); channels/n_heads/head_dim become the `ActivationSpec`."""
+    read: Reader
+    channels: int | None = None  # None only for attention patterns
+    n_heads: int | None = None  # set when the channel axis is heads x head_dim (attention patterns: n_heads)
+    head_dim: int | None = None
+    label: str | None = None
+    description: str | None = None
+
+
+class ArchAdapter:
+    """Base class. Subclass, fill in the class attributes, implement `probe`, `dims` and `acts`, and register with
+    `@register` (see `archs/gpt2.py` for a small complete example)."""
+
+    name: str = ""
+    # `config.model_type` values that select this adapter directly. Adapters are also tried structurally (`probe`),
+    # so a model with a different `model_type` but the same module layout works without being listed here.
+    model_types: tuple[str, ...] = ()
+    layers_path: str = ""  # dotted path from the model to the list of blocks, e.g. "model.layers"
+
+    def probe(self, root: torch.nn.Module) -> list[str]:
+        """Paths this adapter needs but `root` lacks. An empty list means the model is compatible."""
+        raise NotImplementedError
+
+    def dims(self, root: torch.nn.Module) -> Dims:
+        raise NotImplementedError
+
+    def acts(self, root: torch.nn.Module, dims: Dims) -> dict[str, ActDef]:
+        """The activations this model exposes. Keys are registry ids (`capture.REGISTRY`); anything the model
+        does not have (QK-norm, RoPE, a gate, ...) is simply left out. Display order is the registry's."""
+        raise NotImplementedError
+
+    def prepare(self, model, ctx: TraceCtx, act: str) -> None:
+        """Runs inside the trace before any block executes; may stash values in `ctx.aux` for readers.
+        Only do work `act` needs, since every hook adds cost."""
+
+
+# ----- helpers shared by adapters -----
+def get(obj, path: str):
+    """Resolve a dotted path; all-digit parts index. Works on `nn.Module`s and nnsight envoys alike."""
+    for part in path.split("."):
+        obj = obj[int(part)] if part.isdigit() else getattr(obj, part)
+    return obj
+
+
+def has(obj, path: str) -> bool:
+    try:
+        get(obj, path)
+    except (AttributeError, IndexError, TypeError, KeyError):
+        return False
+    return True
+
+
+def missing_paths(root, paths: list[str], block_paths: list[str], layers_path: str) -> list[str]:
+    """Which of `paths` (relative to the model) and `block_paths` (relative to block 0) are absent from `root`."""
+    missing = [p for p in [layers_path, *paths] if not has(root, p)]
+    if missing:
+        return missing
+    return [f"{layers_path}.0.{p}" for p in block_paths if not has(get(root, layers_path)[0], p)]
+
+
+def layer_output(block) -> torch.Tensor:
+    out = block.output
+    return out[0] if isinstance(out, tuple) else out
+
+
+class PreNormBlockAdapter(ArchAdapter):
+    """Sequential pre-norm blocks (`x + attn(norm1(x))`, then `x + mlp(norm2(x))`): Llama, Qwen, Mistral, GPT-2, ...
+
+    Implements the activations that do not depend on how the projections are laid out (residual stream, norms,
+    attention pattern/context/output, MLP output). Subclasses add q/k/v (+ RoPE, QK-norm) and the MLP internals.
+    Module names are relative to a block.
+    """
+
+    norm1: str  # pre-attention norm
+    attn: str
+    o_proj: str  # attention output projection, relative to the block
+    norm2: str  # pre-MLP norm
+    down_proj: str  # MLP output projection, relative to the block
+
+    def shared_acts(self, d: Dims) -> dict[str, ActDef]:
+        n1, n2, o, dn = self.norm1, self.norm2, self.o_proj, self.down_proj
+        D, H, Dh = d.hidden, d.n_heads, d.head_dim
+        o_name, dn_name = o.rsplit(".", 1)[-1], dn.rsplit(".", 1)[-1]  # labels name the module, not its path
+
+        return {
+            "resid_pre": ActDef(lambda b, c: c.tok(get(b, n1).input), D),
+            "resid_mid": ActDef(lambda b, c: c.tok(get(b, n2).input), D),
+            "resid_post": ActDef(lambda b, c: c.tok(layer_output(b)), D),
+            "attn_norm": ActDef(lambda b, c: c.tok(get(b, n1).output), D, label=f"attn_norm — {n1}"),
+            "attn_pattern": ActDef(lambda b, c: c.host(get(b, self.attn).output[1], torch.float16)[0],
+                                   None, H, None),
+            "attn_ctx": ActDef(lambda b, c: c.tok(get(b, o).input), H * Dh, H, Dh, label=f"attn_ctx — {o_name} input"),
+            "o": ActDef(lambda b, c: c.tok(get(b, o).output), D, label=f"o — {o_name} output"),
+            "mlp_norm": ActDef(lambda b, c: c.tok(get(b, n2).output), D, label=f"mlp_norm — {n2}"),
+            "down": ActDef(lambda b, c: c.tok(get(b, dn).output), D, label=f"down — {dn_name} output"),
+        }
+
+    def probe_paths(self) -> tuple[list[str], list[str]]:
+        """(model-level paths, block-level paths) that must exist; subclasses extend."""
+        return [], [self.norm1, self.norm2, self.attn, self.o_proj, self.down_proj]
+
+    def probe(self, root) -> list[str]:
+        model_paths, block_paths = self.probe_paths()
+        return missing_paths(root, model_paths, block_paths, self.layers_path)

actlens/archs/gpt2.py

@@ -0,0 +1,42 @@
+"""GPT-2 (and DistilGPT-2): LayerNorm pre-norm blocks, fused QKV, learned absolute positions (no RoPE),
+non-gated MLP. A compact example of adding an architecture: a layout, the fused-QKV split and the MLP internals."""
+from __future__ import annotations
+
+from .base import ActDef, Dims, PreNormBlockAdapter, get
+from .registry import register
+
+
+@register
+class GPT2Adapter(PreNormBlockAdapter):
+    name = "gpt2"
+    model_types = ("gpt2",)
+    layers_path = "transformer.h"
+    norm1, attn, o_proj = "ln_1", "attn", "attn.c_proj"
+    norm2, down_proj = "ln_2", "mlp.c_proj"
+
+    def probe_paths(self):
+        model_paths, block_paths = super().probe_paths()
+        return ([*model_paths, "transformer.wte", "transformer.wpe"],
+                [*block_paths, "attn.c_attn", "mlp.c_fc", "mlp.c_proj"])
+
+    def dims(self, root) -> Dims:
+        cfg = root.config  # GPT2Config maps hidden_size / num_attention_heads onto n_embd / n_head
+        return Dims(n_layers=len(get(root, self.layers_path)), hidden=cfg.hidden_size,
+                    inter=getattr(cfg, "n_inner", None) or 4 * cfg.hidden_size, n_heads=cfg.num_attention_heads,
+                    n_kv_heads=cfg.num_attention_heads, head_dim=cfg.hidden_size // cfg.num_attention_heads)
+
+    def acts(self, root, d: Dims) -> dict[str, ActDef]:
+        D, H, Dh, I = d.hidden, d.n_heads, d.head_dim, d.inter
+        out = self.shared_acts(d)
+
+        def qkv(i, name):  # c_attn's output is [q | k | v] along the channel axis
+            return ActDef(lambda b, c: c.tok(b.attn.c_attn.output)[:, i * D:(i + 1) * D], D, H, Dh,
+                          label=f"{name} — c_attn[{name}]")
+
+        out |= {
+            "q": qkv(0, "q"), "k": qkv(1, "k"), "v": qkv(2, "v"),
+            "up": ActDef(lambda b, c: c.tok(b.mlp.c_fc.output), I, label="up — c_fc"),
+            # c_proj's input is exactly act(c_fc(x)); reading it avoids depending on how `act` is implemented.
+            "mlp_act": ActDef(lambda b, c: c.tok(b.mlp.c_proj.input), I),
+        }
+        return out

actlens/archs/llama.py

@@ -0,0 +1,79 @@
+"""Llama-style decoders: Llama, Qwen2/2.5/3, Mistral, SmolLM, ... (gated MLP, RoPE, optional QK-norm, GQA)."""
+from __future__ import annotations
+
+import torch
+
+from .base import ActDef, Dims, PreNormBlockAdapter, TraceCtx, get, has
+from .registry import register
+from .rope import apply_rope
+
+ROPE_ACTS = ("q_rope", "k_rope")
+
+
+@register
+class LlamaAdapter(PreNormBlockAdapter):
+    name = "llama"
+    model_types = ("llama", "qwen2", "qwen3", "mistral")
+    layers_path = "model.layers"
+    norm1, attn, o_proj = "input_layernorm", "self_attn", "self_attn.o_proj"
+    norm2, down_proj = "post_attention_layernorm", "mlp.down_proj"
+
+    def probe_paths(self):
+        model_paths, block_paths = super().probe_paths()
+        return ([*model_paths, "model.embed_tokens"],
+                [*block_paths, "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj",
+                 "mlp.gate_proj", "mlp.up_proj", "mlp.act_fn"])
+
+    def dims(self, root) -> Dims:
+        cfg = root.config
+        n_heads = cfg.num_attention_heads
+        return Dims(n_layers=len(get(root, self.layers_path)), hidden=cfg.hidden_size,
+                    inter=cfg.intermediate_size, n_heads=n_heads,
+                    n_kv_heads=getattr(cfg, "num_key_value_heads", None) or n_heads,
+                    head_dim=getattr(cfg, "head_dim", None) or cfg.hidden_size // n_heads)
+
+    def acts(self, root, d: Dims) -> dict[str, ActDef]:
+        block0 = get(root, self.layers_path)[0]
+        has_qk_norm = has(block0, "self_attn.q_norm") and has(block0, "self_attn.k_norm")
+        has_rope = has(root, "model.rotary_emb")
+        act_fn_is_module = isinstance(block0.mlp.act_fn, torch.nn.Module)
+        nH, nKV, Dh, I = d.n_heads, d.n_kv_heads, d.head_dim, d.inter
+        out = self.shared_acts(d)
+
+        def proj(name, heads):
+            return ActDef(lambda b, c: c.tok(get(b, f"self_attn.{name}_proj").output), heads * Dh, heads, Dh)
+
+        def qk_norm(name, heads):
+            return ActDef(lambda b, c: c.tok(get(b, f"self_attn.{name}_norm").output), heads * Dh, heads, Dh)
+
+        def rope(name, heads):
+            src = f"self_attn.{name}_norm" if has_qk_norm else f"self_attn.{name}_proj"
+
+            def read(b, c: TraceCtx):
+                x = c.host(get(b, src).output).reshape(c.n_tokens, heads, Dh)
+                return apply_rope(x, *c.aux["rope"]).reshape(c.n_tokens, -1)
+            return ActDef(read, heads * Dh, heads, Dh)
+
+        out |= {"q": proj("q", nH), "k": proj("k", nKV), "v": proj("v", nKV)}
+        if has_qk_norm:
+            out |= {"q_norm": qk_norm("q", nH), "k_norm": qk_norm("k", nKV)}
+        if has_rope:
+            out |= {"q_rope": rope("q", nH), "k_rope": rope("k", nKV)}
+
+        def silu(b, c):
+            if act_fn_is_module:
+                return c.tok(b.mlp.act_fn.output)
+            return c.tok(torch.nn.functional.silu(b.mlp.gate_proj.output))
+
+        out |= {
+            "gate": ActDef(lambda b, c: c.tok(b.mlp.gate_proj.output), I),
+            "up": ActDef(lambda b, c: c.tok(b.mlp.up_proj.output), I),
+            "silu": ActDef(silu, I),
+            "swiglu": ActDef(lambda b, c: c.tok(b.mlp.down_proj.input), I),
+        }
+        return out
+
+    def prepare(self, model, ctx: TraceCtx, act: str) -> None:
+        if act in ROPE_ACTS:  # cos/sin are produced before the first block runs
+            cos, sin = model.model.rotary_emb.output
+            ctx.aux["rope"] = (ctx.host(cos), ctx.host(sin))