TRF025

Attention masks must be built once in the model, not rebuilt inside a layer or attention module.

   
Default Enabled
Scope Models added on or after 2026-06-20
Source mlinter/trf025.py
Show in terminal mlinter --rule TRF025

What it does

In modeling_*.py and modular_*.py, checks calls to a mask factory (the masking_utils entry points create_causal_mask, create_bidirectional_mask, create_sliding_window_causal_mask, create_chunked_causal_mask, create_masks_for_generate, and any create_*_mask helper) that occur inside a class which does not inherit from PreTrainedModel. Top-level models and sub-models are where masks are meant to be created, so only plain nn.Module blocks — layers, attention modules, encoders — are in scope.

Why is this bad?

Mask construction is O(sequence length squared) work that does not vary per layer. Building it inside the layer repeats that cost once per layer, and because each layer then owns its own mask, the attention backends can no longer be given a single prepared mask — which is how manual padding fixups and per-layer mask divergence get introduced. Create it once in the model and pass it down.

Example

 class AcmeLayer(nn.Module):
     def forward(self, hidden_states, attention_mask=None, **kwargs):
-        attention_mask = create_causal_mask(
-            config=self.config, input_embeds=hidden_states, attention_mask=attention_mask, ...
-        )
         return self.self_attn(hidden_states, attention_mask, **kwargs)

 class AcmeModel(AcmePreTrainedModel):
     def forward(self, input_ids=None, attention_mask=None, **kwargs):
+        causal_mask = create_causal_mask(
+            config=self.config, input_embeds=inputs_embeds, attention_mask=attention_mask, ...
+        )
         for layer in self.layers:
-            hidden_states = layer(hidden_states, attention_mask, **kwargs)
+            hidden_states = layer(hidden_states, causal_mask, **kwargs)

Suppressing this rule

Add a # trf-ignore: TRF025 comment on the flagged line or the line directly above it. See Suppressing rules for whole-file directives and when a suppression is the wrong answer.