fix: keep the causal mask for broadcast 4D padding masks

tests/integration/model_bridge/test_openai_gpt_adapter.py

@@ -1,9 +1,13 @@
 """Integration tests for the OpenAI GPT-1 architecture adapter."""
 
+import copy
+
 import pytest
 import torch
+from transformers import OpenAIGPTConfig, OpenAIGPTLMHeadModel
 
 from transformer_lens.model_bridge.bridge import TransformerBridge
+from transformer_lens.model_bridge.sources import build_bridge_from_module
 
 MODEL = "openai-community/openai-gpt"
 
@@ -80,3 +84,48 @@ def test_greedy_generation_is_coherent(self, bridge):
         )
         assert isinstance(out, str)
         assert len(out) > len("the capital of france is")
+
+
+@pytest.fixture(scope="module")
+def tiny_models():
+    """A tiny random bridge and an unwrapped copy: the bridge wraps the blocks it is given."""
+    torch.manual_seed(0)
+    config = OpenAIGPTConfig(vocab_size=64, n_embd=32, n_layer=2, n_head=4, n_positions=16)
+    config.architectures = ["OpenAIGPTLMHeadModel"]
+    hf_model = OpenAIGPTLMHeadModel(config).eval()
+    reference = copy.deepcopy(hf_model)
+    return build_bridge_from_module(hf_model, "OpenAIGPTLMHeadModel", hf_config=config), reference
+
+
+class TestOpenAIGPTAttentionMask:
+    """GPT-1 hands its blocks a [batch, 1, 1, seq] padding mask, with the causal
+    mask left to the attention module."""
+
+    PAD = 63
+
+    def test_right_padded_batch_matches_hf(self, tiny_models):
+        bridge, reference = tiny_models
+        tokens = torch.tensor([[11, 12, 13, 14], [21, 22, self.PAD, self.PAD]])
+        mask = (tokens != self.PAD).long()
+        with torch.no_grad():
+            logits = bridge(tokens, attention_mask=mask, prepend_bos=False)
+            expected = reference(tokens, attention_mask=mask).logits
+        assert torch.allclose(logits, expected, atol=1e-5)
+
+    def test_left_padded_row_matches_the_row_alone(self, tiny_models):
+        bridge, reference = tiny_models
+        tokens = torch.tensor([[11, 12, 13, 14], [self.PAD, self.PAD, 21, 22]])
+        with torch.no_grad():
+            logits = bridge(tokens, attention_mask=(tokens != self.PAD).long(), prepend_bos=False)
+            expected = reference(torch.tensor([[21, 22]])).logits
+        assert torch.allclose(logits[1, 2:], expected[0], atol=1e-5)
+
+    def test_mask_does_not_open_attention_to_later_tokens(self, tiny_models):
+        bridge, _ = tiny_models
+        tokens = torch.tensor([[11, 12, 13, 14]])
+        with torch.no_grad():
+            _, cache = bridge.run_with_cache(
+                tokens, attention_mask=torch.ones_like(tokens), prepend_bos=False
+            )
+        pattern = cache["blocks.0.attn.hook_pattern"]
+        assert torch.all(pattern.triu(diagonal=1) == 0)

tests/unit/model_bridge/generalized_components/test_mask_slicing.py

@@ -168,3 +168,27 @@ def test_unresolvable_per_layer_geometry_stays_visibly_flat(self) -> None:
     def test_module_head_dim_resolves_when_index_missing(self) -> None:
         shapes = self._hook_shapes(width=16, layer_idx=None, head_dim=4)
         assert shapes == {"q": (2, 5, 4, 4), "z": (2, 5, 4, 4)}, shapes
+
+
+def test_broadcast_padding_mask_keeps_the_causal_mask() -> None:
+    """GPT-1 passes [batch, 1, 1, kv]: 4D, but padding only."""
+    bridge = _attn_bridge()
+    kv_len, n_heads = 4, 2
+    mask = torch.zeros(1, 1, 1, kv_len)
+    scores = torch.zeros(1, n_heads, kv_len, kv_len)
+
+    out = bridge._apply_reconstruct_attention_mask(scores, mask, seq_len=kv_len)
+    future = ~torch.tril(torch.ones(kv_len, kv_len, dtype=torch.bool))
+    assert torch.all(out[..., future] < 0), out
+    assert torch.all(out[..., ~future] == 0), out
+
+
+def test_single_query_row_mask_on_cached_decode_is_not_over_masked() -> None:
+    """One query row is also what a causal 4D mask looks like on a decode step."""
+    bridge = _attn_bridge()
+    kv_len, n_heads = 4, 2
+    mask = torch.zeros(1, 1, 1, kv_len)
+    scores = torch.zeros(1, n_heads, 1, kv_len)
+
+    out = bridge._apply_reconstruct_attention_mask(scores, mask, seq_len=kv_len, q_seq_len=1)
+    assert torch.all(out == 0), out

transformer_lens/model_bridge/generalized_components/attention.py

@@ -585,8 +585,9 @@ def _apply_reconstruct_attention_mask(
         """Apply causal and optional attention masking to reconstructed scores.
 
         HuggingFace-style 4D masks already encode causal semantics, so they are
-        treated as authoritative. Lower-rank masks do not, so the local causal
-        mask is still applied before adding the caller-provided padding mask.
+        treated as authoritative. Lower-rank masks and 4D masks broadcast over the
+        query rows ([batch, 1, 1, kv_seq_len], as GPT-1 passes) do not, so the local
+        causal mask is still applied before adding the caller-provided padding mask.
 
         Args:
             attn_scores: Attention scores [batch, heads, q_seq_len, kv_seq_len].
@@ -599,7 +600,11 @@ def _apply_reconstruct_attention_mask(
             q_seq_len = seq_len
         min_dtype = torch.finfo(attn_scores.dtype).min
         mask_value = -torch.inf if self.compatibility_mode else min_dtype
-        use_direct_hf_mask = attention_mask is not None and attention_mask.ndim >= 4
+        use_direct_hf_mask = (
+            attention_mask is not None
+            and attention_mask.ndim >= 4
+            and attention_mask.shape[-2] != 1
+        )
         # Bidirectional attention (encoders) and cross-attention have no causal
         # structure, so only synthesize the triangular mask for causal self-attention.
         apply_causal = self.is_causal and not self.is_cross_attention