[Proposal] KV-group destinations for direct path patching on GQA/MQA models

Open
  • Is something wrong?
#1,954 1 comment 0 reactions 1 assignee View on GitHub

Maintainers usually reply within 1 day

@emerardd is already working on this.

Since Oct 29, 2026.

Assessment

This issue has not been assessed yet.

Description

complexity-moderate enhancement TransformerBridge
Proposal

Extend the existing direct-path patching functions with an explicit KV-group destination mode. This would measure a source attention head's path through one later layer's shared K or V group, using the bridge's existing pre-norm input forks.

Motivation

Ordinary K/V activation patching already operates on KV heads, and direct Q-path patching works on GQA models. Direct K/V-path patching is deliberately rejected because its destinations and result axes are defined per query head (#1901).

That leaves a specific question without a high-level sweep: how does one source head affect the output through a particular shared K or V input? Restoring the whole clean K/V activation also restores changes from other upstream components. A KV-group destination would isolate the selected source's contribution while respecting the model's native sharing structure.

Pitch

Keep the current behavior as the default, and add an opt-in destination unit, provisionally destination_unit="kv_group", for component="k" or "v". Source indices remain query/output-head indices. Single-source results would be [n_layers, n_kv_heads]; all-sources results would retain query-head source axes and use a KV-group destination axis.

For source head (s, h), compute its clean-minus-corrupted residual write, delta = (z_clean - z_corrupt) @ W_O. Add this only to the selected destination group's pre-norm residual copy, then apply the destination norm and K/V projection normally. With four Q heads and two KV heads, intervening on group 0 reaches query heads 0 and 1 through that shared branch.

The initial scope would be standard layouts with a uniform destination KV count and existing split-QKV forks. Current source-architecture restrictions would remain; heterogeneous KV counts, post-repeat per-query-head interventions, and EAP/faithfulness extensions would be separate scope decisions. On MHA, a group has one head and should reproduce the existing K/V results.

Proposed delivery:

  • One PR updating the single-source and all-sources functions, tests, and documentation. Tests would cover MHA equivalence, MQA/GQA branch isolation against an independent projection-level oracle, no-op interventions, unsupported layouts, and flag restoration on success and exceptions. Include a fixed-checkpoint, multi-prompt GQA example with clean/corrupt baselines and explicit group labels.

Two points to agree on first: would you prefer an opt-in argument or a separate KV-group entry point? Also, the fork currently recomputes the raw norm and can bypass existing ln1.hook_out edits. Should that composition be addressed before this extension, or explicitly rejected when unsupported?

Alternatives

Researchers can implement the intervention with hooks today; the extension would centralize group indexing, sweeps, and result semantics in the existing API. Editing one K/V copy after repeat_kv is another possible experiment, but it breaks native sharing and answers a different question.

Additional context

I checked feasibility on dev at 11197c23. Across 14 config-built Llama interventions covering MQA, GQA, and MHA, the input-fork intervention matched an independent projection-output oracle exactly. On cached HuggingFaceTB/SmolLM2-135M at revision 93efa2f097d58c2a74874c7e644dbc9b0cee75a2, both K and V interventions also matched the oracle, with no change to the destination layer's out-of-group head outputs. These checks establish the intervention mechanics; the proposed multi-prompt example would validate its use on a research task.

References: GQA, IOI/path patching.

Checklist
  • I have checked that there is no similar issue in the repo.
Dominant language
Python
Stars
3.9k
Forks
708
Avg merge
1d 10h
Merged PRs (30d)
82

Getting set up

Open in Codespaces

Starts the project's dev container in your browser, under your own GitHub account.

  • No Dockerfile or Docker Compose file
  • Has a pull request template
  • No contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from TransformerLensOrg/TransformerLens

All issues in TransformerLensOrg/TransformerLens

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.