Manage the input_pos internally in text llm runner

[larryliu0820]

🚀 The feature, motivation and pitch

Instead of exposing input_pos in generate_from_pos() API, we should redesign the API to hide the input_pos argument as an internal state.

We should support these features:

1. generate with an input prompt -> uses the current context, creates the response adds it to context, and adjusts start position of KV caching internally
2. Add context - used to hydrate KV cache for loading historical chat, adjusts start position internally when generate is called after it
3. clear context - remove prefilled tokens and reset start position

To be more specific,

• Add a private field pos_ and manage it in all APIs.
• Keep the generate() API, but instead of assuming a start pos of 0, use the pos_ field.
• Add prefill() API to be able to take chat history.
• Add reset() API to reset pos_ to 0.

Alternatives

No response

Additional context

No response

RFC (Optional)

No response

[nrjpoddar]

Just want to add that generate should internally add the response of the LLM in KV cache. Making calls to Prefill with last response for ongoing conversations is not ideal.

[larryliu0820]

Just want to add that generate should internally add the response of the LLM in KV cache. Making calls to Prefill with last response for ongoing conversations is not ideal.

We are doing that right now. We just don't have internal pos yet

[kimishpatel]

I think we need APIs on modules that allow for getting and setting states such as

ETensor state = module.get_state("transformer.layers.0.attention.k_cache")
module.set_state("transformer.layers.0.attention.k_cache", new_state)

cc: @JacobSzwejbka

[larryliu0820]

I think we need APIs on modules that allow for getting and setting states such as

ETensor state = module.get_state("transformer.layers.0.attention.k_cache")
module.set_state("transformer.layers.0.attention.k_cache", new_state)

cc: @JacobSzwejbka

Seems orthogonal?

[mergennachin]

@nrjpoddar take a look at this PR (#14028)