fix(bridge): expand grouped K/V heads in QK/OV and composition circuits on GQA models

[TravisHaa]

What this fixes

bridge.QK, bridge.OV, and bridge.all_composition_scores() crash on every GQA model (Llama, Qwen, Gemma, Mistral, most MoE). The bridge stacks W_Q/W_O per query head ([n_layers, n_heads, ...]) but keeps W_K/W_V grouped ([n_layers, n_kv_heads, ...]), so FactoredMatrix hits a broadcast error on the leading head axes before any analysis can run. HookedTransformer already expands grouped K/V in its weight properties; the bridge never did.

Fixes #1553.

Changes

• New _expand_kv_heads() in transformer_bridge.py: repeat_interleaves a stacked grouped K/V head axis up to n_heads. Query head h reads kv head h // (n_heads // n_kv_heads) — the same layout GroupedQueryAttention.W_K/W_V expose on HT, and the same routing HF uses at runtime. No-op for MHA and non-4D tensors; raises ValueError if n_heads isn't a multiple of n_kv_heads rather than silently mis-expanding.
• Applied at all seven grouped FactoredMatrix sites: QK, OV, QK_for_attn_layers(), OV_for_attn_layers(), and the three inside all_composition_scores(). The issue lists the first four; composition scores have the identical mismatch (their head labels already assume n_heads per layer), so I fixed them in the same pass rather than leaving GQA composition scores broken.
• Raw bridge.W_K/W_V stacks and per-block weights stay grouped — only the circuit views expand.
• Removed the inert ungroup_grouped_query_attention kwarg from TransformerBridgeConfig. No bridge code reads it, and the only tests exercising the flag build HookedTransformer objects; HT's flag and set_ungroup_grouped_query_attention are untouched.

Tests

• tests/unit/model_bridge/test_expand_kv_heads.py (new, synthetic): the head mapping itself, MHA/non-4D pass-through as the identical tensor, indivisible-count ValueError.
• tests/integration/model_bridge/test_attention_weight_accessors.py (extended, reuses the file's tiny random GQA Llama fixture — no Hub access): QK/OV factors align per query head against the per-block grouped weights, both *_for_attn_layers variants, composition scores in all three modes, raw storage stays grouped, and an MHA bit-identity control on tiny GPT-2.
• tests/integration/model_bridge/test_bridge_qk_ov_vs_hooked_gqa.py (new, @pytest.mark.slow since Qwen2-0.5B isn't CI-cached): bridge QK/OV factors, the layer-0 AB product, and the *_for_attn_layers variants match HookedTransformer.from_pretrained_no_processing via torch.testing.assert_close.

Acceptance criteria from #1553

• bridge.QK/OV match HookedTransformer.QK/OV within fp tolerance on a small GQA model — exact on Qwen2-0.5B locally (max |Δ| = 0.0 on all four factors)
• QK_for_attn_layers/OV_for_attn_layers return the same ungrouped alignment
• Grouped-storage tests unchanged (raw grouped storage additionally asserted in the new tests)
• Regression test asserts non-GQA (MHA) QK/OV unchanged
• rg ungroup transformer_lens/model_bridge/ transformer_lens/config/transformer_bridge_config.py returns zero hits
• uv run mypy . clean (421 files); make format applied; patching.py untouched
• make unit-test: 5097 passed / 0 failed / 59 skipped / 10 xfailed locally, with one caveat below

Local caveat: five generation-exercising unit test files SIGBUS on my macOS-arm64 machine. That's pre-existing, not from this diff — it reproduces identically on the clean dev-4.x base under both torch 2.7.1 and 2.11.0, and matches the "Upstream PyTorch/HF on M-series Macs" quarantines in tests/QUARANTINES.md. The 5097 figure is the tier with those five files excluded; CI should be treated as authoritative for the full tier. I also ran test_analysis_methods.py (22 passed) and test_attention_weight_accessors.py (13 passed) directly as regressions for the touched surfaces.

Relation to #1584

Overlaps — I was assigned on the issue and had this in flight. Beyond the overlap, this PR also fixes the three broken sites in all_composition_scores(), adds the HookedTransformer-parity test from the acceptance list, keeps the rg ungroup check at zero hits, and guards indivisible head counts.