Fix hidden state selection on Gemma4 assistant's first prefill step (#48704)
On the first assisted-generation step, n_last_matches is always 0,
so the old code indexed last_hidden_state[:, 0:1], picking the hidden
state of the first prompt token instead of the last one from the main
model's prefill pass. Guard this case with an is_main_model_prefill flag
that selects last_hidden_state[:, -1:] on the first step and falls back
to n_last_matches indexing afterwards.
Co-authored-by: Codex
Co-authored-by: Sanggyu Lee <sg5.lee@samsung.com>