Circuit tracer speedup

Anthropic's circuit-tracing work reverse-engineers a model's computation on a prompt into an attribution graph over a sparse set of features [1, 3]Circuit Tracing: Revealing Computational Graphs in Language ModelsEmmanuel Ameisen, Jack Lindsey, Adam Pearce, Wes Gurnee, et al.transformer-circuits.pub ↗On the Biology of a Large Language ModelJack Lindsey, Wes Gurnee, Emmanuel Ameisen, Brian Chen, et al.transformer-circuits.pub ↗. They also released the open-source circuit-tracer library and an interactive graph viewer.

The method takes a GPT-style decoder-only transformer and, on a single prompt, builds a replacement model: a linear approximation of the original that is far easier to interpret, because the MLP neurons are replaced with a sparse set of interpretable features. Linearizing is possible because the non-linear parts can be frozen on that one prompt. The attention patterns are precomputed and held constant, the layer norm's per-token scale is fixed so it becomes a linear map, and each MLP is replaced by a transcoder that predicts the MLP's output as a sparse sum over the features that fired. What is left is a linear operator on the residual stream, and the edges of the attribution graph are the attributions between its nodes.

Because the replacement model is linear, its backward pass is a transpose, and one backward pass from a target yields that target's entire row of edges. Making that sweep cheap is most of the work, and it is what the task asks for. The agent is given a working attribution library and has to make its backward sweep substantially faster and lighter, on CPU and on GPU, while returning exactly the same graph. Applied at scale the method runs into three problems the agent can attack: VRAM that scales steeply with prompt length, an edge matrix that has to be fully materialized because there is no streaming or sparsity option, and no native support for explaining more than one output token at a time.

The agent works in an isolated container with no general internet access. It is given fastattr, a two-function attribution engine whose starting implementation delegates to a reference tracer vendored alongside it, so the code it inherits is correct and as slow as the reference by construction. The contract it has to keep is written down in attrgraph/SPEC.md, and a development pool of stored reference graphs lets it check correctness locally as often as it likes. Two tools matter most: one checks its engine against those graphs, and one times the engine the way the grader will, on the same grid, so the agent can measure itself before spending a submission. Its workspace is a git repo and only what it commits is submitted.

# Working in this environment

Everything lives under `/app`:

```
/app/attrgraph/       SPEC.md (the contract), dev/ (the development pool: bundles, prompts,
                      reference graphs), README.md (data dictionary)     read-only
/app/tools/           check · bench · submit · clock                     read-only
/app/src/             your workspace: the engine, a git repo             yours
/app/experiments/     history.jsonl, your lab notebook                   yours
/app/submissions/     your submission records, managed by `submit`; never edit by hand
```

The environment is offline except for the leaderboard. A GPU is available and both CPU and CUDA axes are graded. torch,
transformer-lens and the reference tracer's dependencies are preinstalled in `/opt/venv`
(the default `python3`). No runtime `pip install`.

## What is in /app/src

`src/fastattr/` (the contract package, delegating), `src/circuit_tracer/` (the vendored reference implementation, yours to read, borrow from, or restructure; the stored reference graphs do not change if you edit it, so editing it only degrades your own safety net), `tests/` (contract and agreement tests against the vendored reference), `bench/` (the original dev profiler).

The contract is `/app/attrgraph/SPEC.md`: `fastattr.prepare(model)` once per model,
`fastattr.attribute_fast(handle, prompt)` per prompt, returning a graph with
`adjacency_matrix`, `active_features`, `selected_features` that matches the reference
tracer's no-cap graph. Everything under `/app/src` is yours to restructure as long as
that interface holds.

## Models and prompts

Bundles are synthetic transformers with random transcoders (`/app/attrgraph/README.md`).
The graded grid, exactly:

| axis | bundle | layers x width | features/layer | prompt tokens | prompts | device |
|---|---|---|---|---|---|---|
| `small/16`@cpu | small | 2 x 512 | 64 | 16 | 3 | cpu |
| `small/16`@cuda | small | 2 x 512 | 64 | 16 | 3 | cuda |
| `small/48`@cpu | small | 2 x 512 | 64 | 48 | 3 | cpu |
| `small/48`@cuda | small | 2 x 512 | 64 | 48 | 3 | cuda |
| `large/16`@cpu | large | 8 x 512 | 1024 | 16 | 3 | cpu |
| `large/16`@cuda | large | 8 x 512 | 1024 | 16 | 3 | cuda |
| `large/32`@cpu | large | 8 x 512 | 1024 | 32 | 2 | cpu |
| `large/32`@cuda | large | 8 x 512 | 1024 | 32 | 2 | cuda |
| `deep/16`@cpu | deep | 12 x 512 | 1024 | 16 | 3 | cpu |
| `deep/16`@cuda | deep | 12 x 512 | 1024 | 16 | 3 | cuda |

Prompts are random token sequences of the listed lengths; the graded pools use the same
bundle shapes and lengths as the dev pool with different seeds, so nothing keyed on the
dev prompts or the dev weights transfers, but everything algorithmic does.

## Checking correctness (free, unlimited)

```
/app/tools/check                 # every dev prompt vs its stored reference graph
```

Same comparator and thresholds as the grader: active-feature Jaccard, selected-feature
Jaccard and edge correlation all at or above 0.99. A failing check means a failing
submission.

## Measuring speed (free, unlimited)

```
/app/tools/bench --label my-idea > run.log 2>&1 ; grep "^dev_score:" run.log
```

`bench` times `attribute_fast` on each dev axis exactly the way the grader does
(single-threaded on CPU or on the GPU with `--device cuda`; `prepare` timed separately and not scored; per prompt one warm-up call, then the median of three timed calls, or a single timed call when a call takes over 20 s) and
prints graphs per second and peak memory per axis plus `dev_score`, the geometric mean
over axes of `graphs_per_second × min(1, memory_budget / peak_memory)`. It appends every
run to `/app/experiments/history.jsonl`.

## How submissions are scored

The public score is computed the same way on the graded pool:

- per axis: `speed × min(1, budget / peak_memory)`, speed in graphs per second, peak
  memory in bytes (budgets per axis are listed in `/app/attrgraph/dev/budgets.json`);
- per device: the geometric mean over that device's axes;
- the task score is the sum of the per-device scores (cpu + cuda), both shown on the leaderboard.

Gates come first: the package must import, `prepare` must finish within ten minutes,
every graded graph must agree with the reference within the thresholds above, and the
measured process may not spawn descendants. Any failed gate scores 0 for that
submission, with the reason in the receipt.

## Submitting

```
/app/tools/submit --name grid-hash-v3 --notes "cell size 0.5 eps"
```

- Seals `git archive HEAD` of `/app/src` (commit first: uncommitted and ignored files
  are not included), uploads it, and returns immediately with an id.
- Grading runs in the background and takes 2-12 minutes: the slower your
  engine, the longer the grade (the unmodified starting tree is the worst case).
  Receipts can additionally queue behind other grades in progress. Keep working
  while a submission grades; never sit and poll.
  `submit --list` shows every submission's state; `submit --wait s0007` blocks.
- The receipt (`/app/submissions/s0007-*/result.json`) carries the public score, the
  per-axis numbers, and each gate's verdict.
- You have unlimited scored submissions for this run; local checks and benchmarks are unlimited. Invalid and rejected attempts never count against any limit.
- Submit when your local numbers move meaningfully, not after every commit.

## Experiment discipline

`/app/src` is initialized as a git repo. Conventions:

- one experimental idea per commit; the checked-out commit is always your best-so-far;
  `git reset --hard HEAD~1` rejects a failed idea
- `.gitignore` (preseeded) keeps logs, results.tsv and scratch out of version control
- print `dev_score: <float>` as the last line of every evaluation (bench does); append
  every attempt to `/app/src/results.tsv`

## Time and long-running commands

- `/app/tools/clock` prints elapsed and remaining time (a safe lower bound). Check it
  whenever you are deciding what to attempt next.
- Runs longer than a couple of minutes: launch them in the background
  (`nohup /app/tools/bench --label x > run.log 2>&1 &`) and poll the log, or give the
  shell an explicit long timeout.

## Rules

- Keep `/app/src` to the package: environment setup, caches and scratch go elsewhere.
- No forking or worker processes inside `attribute_fast`; no caching keyed on prompt
  content; no editing of the stored reference graphs (they are read-only and are not
  the graded ones anyway).
- The graded pools exist nowhere in this environment.

Every graded graph is the no-cap graph: every active feature is a node, and the edges are the direct effects between features, error nodes, tokens and logits. Scored submissions are unlimited, and local checks and benchmarks are free, so the constraint on the agent is the four-hour clock rather than a submission budget.

Every submission is graded by the same grader, run once against the public pool for the score the agent can see and once against the held-out pool for the score that counts.

Grading is per prompt. For each prompt the grader loads it, makes one untimed warm-up call, then three timed calls and takes the median as that prompt's time; a call slower than twenty seconds is timed once instead of three times. Each call's output is compared against the stored pristine graph for that prompt, and the comparison has to clear three tolerances at once: edge correlation at least 0.99, active-feature Jaccard at least 0.99, and selected-feature Jaccard at least 0.99. This is not a pass the submission takes before it is timed. It runs on every prompt as the prompt is timed, and a single disagreement fails the whole submission rather than that one axis. CPU axes carry a second check: if the measured CPU time falls below three tenths of the wall time, the work is taken to be running outside the metered process and the submission fails the same way.

The timed prompts are grouped into five workload shapes, chosen so that a win on a single axis cannot carry the score: two small grids, two large, and one deep, at different batch widths. Each shape is run on CPU and on GPU, giving ten axes. An axis scores the throughput the engine reached across its prompts, in graphs per second, multiplied by a memory factor. Writing $K$ for the prompts on the axis, $t_k$ for prompt $k$'s median time, $M$ for the peak memory the axis reached and $B$ for that axis's memory budget,

The memory factor never pays for using less than the budget; it only scales the axis down, in proportion, when the budget is exceeded. Peak memory is the CUDA peak on GPU axes and resident set size on CPU axes.

Each device is then summarised by the geometric mean over its five axes, and the task score is the sum of the two,

The geometric mean is what makes a single-axis win worthless: a shape left untouched holds the whole device term down, and no amount of speed elsewhere buys it back. If any gate fails the score is zero outright, as it is for a submission that changes nothing or breaks the engine.

Below we lay out one possible solution to a faster algorithm that the agents could take. This does not mean that this is the best solution (we highly doubt that it is), but it is one possible solution.

We want to optimize simultaneously for both memory and speed. One thing you can do, non-trivially, is to remove the autograd backend. circuit-tracer computes its edges with PyTorch's reverse-mode autodiff through the replacement model, which is a perfectly reasonable default. But the replacement model is linear, and the backward pass of a linear operator is just its transpose, so we do not need a general autograd engine to run it. The attention path for one layer is a short chain of linear maps. Writing $X$ for the residual stream entering the block, $\mathrm{LN}$ for the frozen layer norm and $A$ for the frozen attention pattern, the contribution it adds back is

Every step is linear, so the backward pass is each map transposed, applied in reverse order. Writing $d_{\mathrm{out}}$ for the gradient arriving at the block's output,

Only the last line is not literally a matrix transpose: we run the frozen layer norm in reverse, scaling the gradient and then centering it, the opposite order of the forward pass. The MLP path carries no gradient of its own, since the feature activations are frozen, so those four lines are the entire backward sweep. Writing them by hand removes the engine outright, and more than half of circuit-tracer's device time turned out to be elementwise and reduction kernels that came with autograd rather than edge computation. It also exposes every intermediate to optimization, which is what makes everything below possible.

The memory savings come from the same change. To differentiate through the replacement model, reverse-mode autograd keeps the forward graph's saved activations alive across the backward pass, and that retained set grows with the batch of features, the prompt length, and the depth. Our hand-written pass keeps none of it. It holds one running gradient buffer that it reuses as it iterates down the layers, so the footprint does not grow with depth, plus the frozen weights, and per attributed feature only that feature's encoder and decoder rows. Measured end to end, an extra feature costs about 12 KB, against about 74 KB under autograd.

With the engine gone, the next thing to exploit is causality. A node only ever attributes to nodes below it in the network and at or before it in the sequence, and both bounds let us skip work. Layer causality autograd already gets for free, since gradient does not flow above the layer the node was injected at; we keep it explicit, starting the sweep at the highest layer in the batch and iterating down, and sorting each batch by layer and position so that most batches never run the full depth. Temporal causality is the one circuit-tracer cannot exploit. Decoder attention is causal, so a node at position $T$ has no ancestors after $T$, and the sweep only needs the prefix up to the batch's latest position. We compute that prefix per batch and slice every buffer, every pattern, and every per-feature dot down to it; features past the prefix have provably zero gradient and are never materialized at all. The two compound, so early-layer, early-token batches do almost no work while only the late ones pay close to full cost. These were the largest optimizations we measured.

What remains after that is dominated by two matrix multiplies. The projections through the output and value weights account for roughly 99% of the backward FLOPs, and one structural fact removes a large part of the value side. Qwen3-4B uses grouped-query attention: it has 32 query heads but only 8 key-value heads, so each key-value weight is shared by a group of four query heads.

The catch is that this structure is not visible by the time the backward pass sees the weights. circuit-tracer loads models through TransformerLens [2]TransformerLensNeel Nanda, Joseph Bloomgithub.com ↗, which expands the key-value weights back out to one identical copy per query head before handing them over, so the tensor the backward multiplies through is the full 32-head matrix and the grouping is gone. A naive backward then pushes the gradient through all four copies of every weight. This is worth stating plainly because it is not a property of the model or of the method: it is a property of the library the model is loaded through, and it would not show up in the maths.

It can be recovered. The forward pass copies each weight across its group, and the transpose of a copy is a sum, so instead of multiplying every query head's gradient through its own copy, the backward sums the gradient within each group first and multiplies once per key-value head:

where $d_V[h]$ is query head $h$'s slice of the gradient, $W_V[h]$ is its copy of the weight, and $W_V[g]$ is the one weight all of group $g$'s copies are copies of. This is exact, since we are only regrouping a sum we were already computing, and it drops the two dominant projections to roughly 0.625 of their original combined cost on this model. It generalizes to any grouped-query model, and we only collapse the expanded weight after checking that every group really is a bit-identical copy of its first head, so multi-head models keep the old path untouched.

Every one of these is mathematically exact, which is the part that matters here. Measured directly on the sweep prompts, the two backends select identical feature sets at every length and the full adjacency matrices agree with correlation 0.9999; the only disagreements are a handful of near-zero edges that one reduction order rounds to zero and another does not. None of the speedup comes from changing the answer.

On Qwen3-4B the backward sweep ends up roughly 6x faster across the measured range, 6.9x at 5 tokens down to 5.8x at 49, and roughly 6x lighter per attributed feature. The memory difference is the one that decides whether a long prompt runs at all: the baseline grows by about half a gigabyte across the whole sweep while circuit-tracer grows by almost three and a half, and runs out of VRAM not far past the end of it. There is not much room left, either. The optimized matrix multiplies run at about 98% of the card's bf16 tensor-core peak, so what remains is the edge-matrix copy and launch overhead rather than arithmetic, and the next improvements would be structural rather than faster math.