fix(bridge): gate batched-list position_ids on the target model

tests/integration/model_bridge/test_llada_adapter.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import copy
+import functools
 import gc
 import math
 import weakref
@@ -663,6 +664,54 @@ def test_left_padding_does_not_inject_unsupported_position_ids(models: TinyModel
     torch.testing.assert_close(bridge_logits, reference_logits, rtol=1e-5, atol=1e-6)
 
 
+def test_batched_list_input_does_not_inject_unsupported_position_ids() -> None:
+    """Batched list input builds its own attention_mask and position_ids so pad
+    tokens don't contaminate the forward (#1626). The mask is safe for any model;
+    the position_ids are not, and this forward takes neither them nor **kwargs.
+
+    A local bridge rather than the module fixture: this needs a tokenizer, and
+    attaching one to the shared instance would leak into the other tests. The
+    tokenizer is given a BOS so the path under test is reached independently of
+    BOS handling elsewhere.
+    """
+    local = _build_models()
+    tokenizer = _offline_tokenizer()
+    tokenizer.bos_token = "<bos>"
+    local.bridge.tokenizer = tokenizer
+
+    seen: dict = {}
+    original = local.bridge.original_model.forward
+
+    # functools.wraps so inspect.signature() still resolves to the real forward:
+    # the gate reads that signature, and a bare (*args, **kwargs) spy would look
+    # like it accepts position_ids and defeat the check under test.
+    @functools.wraps(original)
+    def _spy(*args, **kwargs):
+        seen.clear()
+        seen.update(kwargs)
+        return original(*args, **kwargs)
+
+    local.bridge.original_model.forward = _spy
+    try:
+        with torch.inference_mode():
+            logits = local.bridge(["token_5 token_7 token_9", "token_5"], return_type="logits")
+        batched = dict(seen)
+        with torch.inference_mode():
+            local.bridge("token_5 token_7 token_9", return_type="logits")
+        unbatched = dict(seen)
+    finally:
+        local.bridge.original_model.forward = original
+
+    assert logits.shape[0] == 2
+    assert "position_ids" not in batched
+    # The mask is still supplied — withholding it would reintroduce the padding
+    # contamination this branch exists to prevent.
+    assert "attention_mask" in batched
+    # Control: a single unbatched string never reached this branch, so the gate
+    # must not have changed anything for it either.
+    assert "position_ids" not in unbatched
+
+
 def test_run_with_cache_exposes_hooks_without_hf_output_attentions(
     models: TinyModels,
 ) -> None:

transformer_lens/model_bridge/transformer_bridge.py

@@ -1734,7 +1734,11 @@ def forward(
                     ).to(self.cfg.device)
                 finally:
                     self.tokenizer.padding_side = _prev_side
-                if "position_ids" not in kwargs:
+                # Gated on the target for the same reason the derivation below is:
+                # a fixed-signature forward raises TypeError on the kwarg, and a
+                # model that owns its own position derivation is overridden by it
+                # (#1626).
+                if "position_ids" not in kwargs and self._accepts_derived_position_ids():
                     position_ids = attention_mask.long().cumsum(-1) - 1
                     position_ids.masked_fill_(attention_mask == 0, 1)
                     kwargs["position_ids"] = position_ids
@@ -2142,9 +2146,12 @@ def _generate_tokens(
                         ).to(self.cfg.device)
                         self.tokenizer.padding_side = _prev_side
                         forward_kwargs["attention_mask"] = attn_mask
-                        position_ids = attn_mask.long().cumsum(-1) - 1
-                        position_ids.masked_fill_(attn_mask == 0, 1)
-                        forward_kwargs["position_ids"] = position_ids
+                        # Same target gate as the forward() path: the mask is safe
+                        # for every model, the derived positions are not (#1626).
+                        if self._accepts_derived_position_ids():
+                            position_ids = attn_mask.long().cumsum(-1) - 1
+                            position_ids.masked_fill_(attn_mask == 0, 1)
+                            forward_kwargs["position_ids"] = position_ids
                     if gen_step_idx == 0:
                         if pixel_values is not None:
                             forward_kwargs["pixel_values"] = pixel_values