Support 2D bool mask for logits_to_keep, expose last_hidden_state

Feature request

Two changes to Gemma4ForConditionalGeneration.forward:

1. Accept a bool tensor for logits_to_keep to support non-contiguous masking.
2. Add a flag independent of output_hidden_states that returns last_hidden_state.

Motivation

SFT training with sequence packing concatenates multiple examples into one sequence, and some spans within that sequence need to be excluded from logits computation. These spans are non-contiguous, so a slice or an int can't express them. Right now this requires monkey-patching forward() to slice hidden_states directly.

last_hidden_state is unexposed for the same reason. Getting the final layer's hidden state for anything other than logits currently requires output_hidden_states=True, which stores every layer's hidden state, then indexing hidden_states[-1]. This keeps unneeded intermediate layers in memory, and forward() still doesn't return last_hidden_state directly, so this also needs a monkey patch.

Both problems come down to the same thing: neither is solvable without a monkey patch.

Your contribution

Willing to submit a PR implementing this for Gemma4ForConditionalGeneration.