Accessing Attention on MLX

[tburleyinfo]

“How can intermediate attention activations (e.g., attention heads or attention weights) be accessed or inspected when running transformer models with MLX?”

[tk1475]

MLX doesn't have PyTorch-style forward hooks, but since models are plain Python objects you can reach intermediates directly.

Hidden states / activations: any nn.Module is callable, so run inputs up to a layer (model.model.layers[i]) or wrap a module's __call__ to stash its output.

Attention weights specifically are trickier: mlx-lm's attention uses the fused mx.fast.scaled_dot_product_attention, which never materializes the full softmax(QKᵀ) matrix, so there's nothing to read out. To inspect them, replace that fused call in the attention module with the manual path:

scores = mx.softmax((q @ k.swapaxes(-1, -2)) * scale + mask, axis=-1)
out = scores @ v   # stash `scores` before this line

Grab the module (model.model.layers[i].self_attn), monkeypatch or subclass its __call__ to compute and keep scores, run a forward, then mx.eval() the captured array before inspecting (MLX is lazy). It gives up the memory and speed the fused kernel saves, so only do it on the layers you're probing. For interpretability across all layers, a small wrapper that reimplements attention to also return scores is cleaner than patching in place.