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.
