[Feature]: guided decoding on TPU

[carlesoctav]

🚀 The feature, motivation and pitch

I’m not sure if this is possible, but right now the execute_model function on the TPUModelRunner is only outputting the predicted token_ids, rather than the distribution of tokens that we can sample from with some guidance (e.g., using outlines). I believe structured output is becoming more common, and most projects that require LLMs need this structured output feature.

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[carlesoctav]

I'm down to work on these features :).

[robertgshaw2-redhat]

Thanks @carlesoctav - to get started, we need to enable support for LogitsProcessing on TPUs. Do you need a pointer to get started?

[bvrockwell]

Thanks so much for lending a hand @carlesoctav ! Indeed, this is super important :)

[carlesoctav]

hi, I've been working on making these features viable and concluded with this approach:

1. Extract the logits_processors in the prepare_sample function (previously outputting just n, p, t params).
2. Pass the logits_processors as one of the parameters for the ModelWrapper.forward class.
3. Iteratively apply the logits_processor (similar to the _apply_logits_processor function).

However, there are still some missing parameters needed for a LogitProcessor, mainly prompt_token_ids or past_tokens_ids and it required sample_indices to extract those params.
may I know how I can get these parameters? Do I need to extract them in the prepare_input function?

also here's the diff for the changes I made:
carlesoctav@34703fc

[bvrockwell]

cc @dyli-google 👍

[bvrockwell]

@Chenyaaang

[dyli-google]

@carlesoctav

Sorry for the late reply.

Is there any update on this? Is carlesoctav@34703fc still the latest commit?

Also, do you want to create a pull request for this?

Thanks.

[Chenyaaang]

I did some investigation yesterday, carlesoctav's way is workable, but given the pr to support structured decoding on GPU V1 (#12388), we only need to do the same thing on v1/worker/tpu_model_runner as v1/worker/gpu_model_runner. I can implement it after the pr is merged.

[russellb]

I believe this was completed by #16499