[emerardd]

Description

Residual decomposition with mlp_input=True assumes that the current layer's attention output feeds its MLP. That is correct for sequential blocks, but parallel blocks feed both branches from resid_pre. On a tiny random NeoX with nonzero attention output biases, the existing component decomposition differs from the actual ln2 input by about 0.4; normalized attribution differs by about 23. Full head/neuron decomposition also includes the nonexistent current-layer attention contribution.

This change declares the MLP-input topology on BlockBridge / ParallelBlockBridge and uses the target block's topology in the analysis paths:

• Component and full residual decompositions include current-layer attention/head writes only for sequential blocks.
• accumulated_bias() excludes the current attention output bias for parallel MLP inputs, while retaining all preceding contributions.
• Accumulated residuals omit nonexistent parallel mid streams.
• Neuron remainders target parallel resid_pre rather than an absent resid_mid, including the fused normalized projection path.

Model forward computation, hook naming, neuron selection and ordinary attention-input decomposition are unchanged. The existing method docstrings now describe parallel MLP-input semantics.

Type of change

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

Validation

The initial real-config NeoX regression had 16 failed / 16 passed before the fix: the parallel cases failed and the sequential controls passed. These models are built from config and do not download weights.

Expanded coverage checks layers 0/1, raw/normalized sums, batched/batchless caches, expanded/unexpanded neurons, vector/matrix projections, bias inclusion and remainder targets. Cross-architecture checks compare decomposition sums to captured normalization inputs, including shared-ln1 and separate-ln2 layouts.

• NeoX regression, mixed-norm decomposition and two existing activation-cache test files: 362 passed.
• Final parallel/sequential architecture identities and expanded NeoX regression: 108 passed (overlaps the NeoX cases above).
• Full mypy .: 389 source files passed.
• Pinned pycln/isort/Black checks on all five changed files and git diff --check: passed.

The complete unit suite, pretrained checkpoints and GPU/MPS hardware were not tested.

Checklist

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation (API docstrings)
• 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 (affected tests run; full suite not run)
• I have not rewritten tests relating to key interfaces which would affect backward compatibility