attention

Masked softmax attention heatmap for one example and query head.

Combine this unregistered template with your trainer (analysis first)::

@analysis(
    "gpt/analyze/attention"
)
class GPTAttention(
    AttentionHeatmapAnalysis,
    GPTPretrain,
):
    pass

The trainer supplies MODEL, DATASET, and restoration/surgery. CONFIG defaults to AttentionAnalysisConfig, extending BaseTrainerConfig. For a custom trainer config, inherit both schemas in a dataclass and set CONFIG on the concrete analysis. This adds analysis/layer, sample, head, max_tokens to its configuration. Layer indices follow discovery order. Execute with base=checkpoint through the job lifecycle.

This analyzes full-sequence batches, not KV-cache decoding. All hosts compute; only the bounded attention matrix is gathered, then host zero creates a CPU plot. Softmax uses all keys before cropping: displayed rows can sum to less than one when undisplayed keys receive attention. An entirely masked row is shown as zero.