Hello team,

I'm currently exploring solutions to serve a custom model and would appreciate your insights on whether my use case is feasible with text-generation-inference (TGI).

My model requires a custom embedding logic based not only on input_ids, but also on their position within the full sequence. More specifically, I’m working on a mesh generation task, and in addition to the tokens, I use two types of position IDs:

When creating the embedding for an input_id, I need to know its position within the full input_ids sequence, including both the prompt and the generated tokens.

I implemented this in a Hugging Face model by customizing the prepare_inputs_for_generation method, ensuring that during generation, I always had access to the full list of input_ids. This allowed me to generate embeddings appropriately.

Now, my question is: How feasible is it to replicate this logic using TGI?

I initially tried vLLM, but found it challenging to access the full sequence of input_ids during the forward pass, particularly with continuous batching where inputs are interleaved.

Any guidance on whether TGI supports this type of positional logic—or if there’s a recommended way to achieve it—would be greatly appreciated!

Thanks in advance!