fix(bridge): accept attention_mask in generate() for pre-padded prompts

[sohv]

generate() had no way to be told which prompt tokens are padding, so an already-padded token tensor generated as though its pads were real context. I found the old behaviour was stranger than I described in the issue: an attention_mask passed by a hopeful caller was absorbed into **multimodal_kwargs, which are merged into the forward kwargs on step 0 only. The first token therefore came out right and every later one drifted, which is harder to spot than a uniform failure.

What I changed

attention_mask is now an explicit parameter, extended by one attended column per generated token, so every step sees a mask spanning the prompt plus what has been generated so far. I also fixed the cached step, which pinned position_ids to total_len - 1 — that counts pad slots, so it is wrong for a left-padded prompt. It now reads the new token's position off the mask.

Separately, padding_side is now applied to token input. generate() has always documented the argument but only used it when tokenizing string or list input, leaving it inert for a tensor. It is consulted only when explicitly passed, because deriving a mask on the default path would change behaviour for every existing caller and would require a real tokenizer where none is needed today. That covers the common single-edge case; only an explicit mask can express an interior gap or a pad id that also occurs as a real token.

Verification

Correct across 3 architectures × 3 padding amounts × 2 cache settings, via both routes:

unpadded                  [262, 3139, 286, 262, 4141, 2066]
padded, no mask           [262, 3139, 286, 262, 1578, 1829]
padded, attention_mask    [262, 3139, 286, 262, 4141, 2066]
padded, padding_side=left [262, 3139, 286, 262, 4141, 2066]

The no-mask result is unchanged from before this PR, so nothing moves for callers who don't opt in.

17 tests in a new file, 11 red on unmodified code. One of them spies on the forward and asserts the mask's width and masked region at every step, which is the part the old step-0 leak got wrong.

Full unit + integration + acceptance: 6449 passed, 1 failed. The failure is test_bridge_hooked_parity_multi_step_optimization, which produces the identical 0.100004 exceeds threshold 0.100000 on unmodified code. I also re-ran the areas closest to this change on their own: all generation suites (129 passed) and every multimodal adapter (35 passed). mypy clean.

Two notes

I did not touch HookedTransformer, per your note. Nothing here depends on it — the accurate framing is just that the bridge's own padding_side argument was inert for token input.

I left **multimodal_kwargs alone for now, for the reason you raised. GPT-2's forward takes **kwargs, so a signature-based check never fires, and any whitelist risks trimming kwargs that should pass through. Worth its own investigation. Making attention_mask explicit already removes the specific silent-swallow from the issue.

One thing that caught me while testing, and which shaped the final design: the Florence-2 benchmark passes a processor's outputs straight through as **extra, including attention_mask. My first version rejected a mask on encoder-decoder paths, and that NotImplementedError landed inside a bare except Exception: continue — surfacing as a misleading "image processor/PIL unavailable" skip rather than an error. On the inputs_embeds and encoder-decoder paths the mask now flows through to the model untouched, exactly as it did via **multimodal_kwargs.

[jlarson4]

Great work on this @sohv! I have just a couple small comments below, otherwise looks good to me.

[jlarson4]

The heuristic requires isinstance(self.tokenizer, PreTrainedTokenizerBase) and a pad id, so on a boot_native bridge (tokenizer None) an explicit generate(tokens, padding_side="left") silently changes nothing. That's the same failure this PR fixes for real tokenizers, on a separate bridge. Since the gate only fires when the caller explicitly passed padding_side, could you raise a legible ValueError when it can't be honored, pointing at attention_mask as the alternative?

[jlarson4]

On those paths the mask is routed back through multimodal_kwargs untouched. This is a behavior you chose for Florence-2-style processor passthrough. "Not supported" reads as "will raise". Could you reword to something like "on encoder-decoder and inputs_embeds paths the mask is forwarded to the model as-is (not grown per step)" so the docstring matches the code and the test?

[sohv]

Both fixed in a81a323.

On the first one, you were right and it's a sharper version of the bug: on a bridge with no tokenizer the argument was inert, so padding_side="left" would have quietly produced the wrong continuation rather than doing nothing visible. Since the heuristic only runs when the caller explicitly passed it, generate() now raises instead, with separate messages for a missing tokenizer and a missing pad_token_id, each pointing at attention_mask. Your exact scenario:

tokenizer: None
ValueError: generate(padding_side=...) reads the padding off the pad token, which
needs a tokenizer; this bridge has none. Pass attention_mask=... to state the
padding directly instead.

I also checked that the alternative the message advertises actually holds rather than assuming it: a left-padded boot_native prompt with an explicit mask generates [1, 10], identical to the unpadded prompt.

On the second — that wording was left over from the version that raised on those paths, and I never updated it when I switched to pass-through. It now reads "On the encoder-decoder and inputs_embeds paths the mask is forwarded to the model as-is rather than grown per step, which is what processors emitting one alongside pixel_values expect", which matches both the code and the test.

Three tests added, 20 in the file. Full unit + integration + acceptance: 6452 passed, with the same pre-existing test_bridge_hooked_parity_multi_step_optimization failure as before. mypy clean.

[jlarson4]

Great stuff @sohv, solid resolution, thanks for getting all that updated! Approved and merging now