[flyingdoubleG]

Description

On JetMoE, per-head attention results and every API built on them raise:

cache.stack_head_results()               # KeyError: 'blocks.0.attn.hook_z'
cache.get_full_resid_decomposition()     # KeyError: 'blocks.0.attn.hook_z'
bridge.set_use_attn_result(True)         # NotImplementedError: ... Per-head result computation is unavailable.

JetMoE's attention is a mixture of attention heads (JetMoeMoA). Each token routes to top_k attention experts, and each expert has its own query and output projection:

# JetMoeAttention.forward
query_states, router_logits, topo_info = self.experts.map(hidden_states)  # [b, s, top_k, n_kv * d_head]
...                                                                        # attention over top_k * n_kv heads
attn_output = self.experts.reduce(attn_output, topo_info)  # per-expert output_linear, * gate, summed, + bias

Expert slot k owns heads k * n_kv .. (k + 1) * n_kv - 1. Its output projection and gate depend on the token's routing. There is no o_proj for hook_z to alias, and no static W_O for ActivationCache to multiply hook_z by.

Fix.

• Adapter (_JetMoeAttentionBridge).
  • hook_z is a real hook point on the input of experts.reduce, [batch, pos, n_heads, d_head]. It is patchable: the reduce call reads the hooked value.
  • Per-head results: for each head, gate · W_out[expert][:, head] @ z_head, with each slot's expert and gate from the attention router.
  • head_results_from_cache computes them from the cached hook_z and the router logits (attn.experts.router.hook_out), with the same top-k + softmax as JetMoeTopKGating.
  • With use_attn_result, the wrapped reduce computes them from HF's routing topology, fires hook_result, and returns their sum plus the MoA bias. supports_attn_result = True.
  • b_O is the MoA bias, so the decompositions' bias row includes it.
• ActivationCache._ensure_head_results. If an attention bridge defines head_results_from_cache, use it instead of hook_z @ W_O. No other bridge defines it.

On a tiny random JetMoE (3 layers, 4 attention experts, top-2, non-zero MoA bias), max abs errors:

	dev	this PR
each layer's heads + b_O vs hook_attn_out	KeyError	6e-8
stack_head_results(incl_remainder=True) vs resid_post	KeyError	6e-8
get_full_resid_decomposition(expand_neurons=False)	KeyError	2e-7
logits with use_attn_result vs HF	NotImplementedError	3e-8

Logits are unchanged without use_attn_result (still exactly HF's), and greedy generate matches HF. expand_neurons=True still fails on JetMoE's MoE MLP, as on the other MoEBridge models.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• I have commented my code, particularly in hard-to-understand areas

• I have made corresponding changes to the documentation

• My changes generate no new warnings

• I have added tests that prove my fix is effective or that my feature works

• New and existing unit tests pass locally with my changes

• I have not rewritten tests relating to key interfaces which would affect backward compatibility

• New tests/integration/model_bridge/test_jetmoe_head_results.py (tiny random JetMoeForCausalLM, built locally) checks that:

  • logits match HF;
  • hook_z has the per-head shape;
  • zeroing one head through hook_z changes the output, and an identity hook does not;
  • each layer's heads plus b_O sum to hook_attn_out;
  • heads plus remainder sum to resid_post;
  • get_full_resid_decomposition sums to its target;
  • with use_attn_result, logits match HF and hook_result equals the cache-computed results;
  • greedy generate matches HF.

• test_jetmoe_adapter.py: hook_z and supports_attn_result on the MoA bridge. A second test checks that _slot_routing recovers each slot's top-k expert and gate from HF's expert-sorted routing topology, using a real JetMoeTopKGating.

On dev, 9 of these tests fail.

Local runs (Windows, CPU, torch 2.9.1, transformers 5.20.2, --jaxtyping-packages=transformer_lens,beartype.beartype, PYTHONUTF8=1):

• New tests and test_jetmoe_adapter.py: 16 passed.
• tests/unit -m "not slow": 8175 passed, 14 failed. The same 14 fail on dev in my environment (ten Gemma 3n adapter tests need the optional timm, three test_sparse_probing.py tests check optimizer stop reasons, and test_faithfulness_recovery_is_graded_not_binary is [Bug Report] test_faithfulness_recovery_is_graded_not_binary passes or fails on the sign of float32 rounding #1993).
• tests/integration/model_bridge -k "resid or decompos or head_result or attn_result or post_norm": 388 passed. The 2 failures also fail on dev here: test_jamba_adapter.py::test_uses_post_norm_B_C (no mixer.conv1d hooks in my transformers version) and an Ouro test whose remote code calls create_causal_mask(input_embeds=...). The Gemma 3n errors need timm. Three Raven tests errored on downloading the checkpoint and pass on a rerun.

black (the locked 23.12.1), isort --check-only and mypy on the changed files pass. In my environment mypy reports three errors in other, unchanged files (_model_state.py, moe.py, olmo_hybrid.py); they also appear on dev.

This was prepared with AI assistance (Claude); I reviewed and tested the change.