Skipping generation for useless tokens, and modiying cacheids about transformers-neuronx HOT 3 CLOSED

Hi @enochlev,

It looks like you are trying to generate the remaining tokens in a sequence given a common prefix. Given this is your goal, you need to ensure that the suffix you are attempting to encode, begins with the correct indices into the KV cache. This should be the last token index of the common prefix. Note that it may still be more performant to generate a larger context encoding depending on your use case

from transformers-neuronx.

Based on your comment, are you are telling me that the cache_ids have to be correctly incrementally aligned? I thought I did so in my example

First forward call: (prefix)
input_ids: [[0, 0, 0, 0, 0, 0, 0 ...., 13, 4013, 5828, 338, 1048, 29871]]
cache_ids: [ 0, 1, 2, 3, 4, 5, 6....., 794, 795, 796, 797, 798, 799]

Standard forward call with live cache
input_ids: [[1058]]
cache_ids: [800]
**probabilities seems to be fine with this

Next forward call to skip 4 forward calls..#what I am trying to implement
input_ids: [[29871, 263, 5381, 767, 1058]]
cache_ids: [801, 802, 803, 804, 805].
**probabilities seemed to get messed up with this even though the cache_ids are aligned

When you are mentioned prefix, are you suggesting the set_prefixed function call?

from transformers-neuronx.

I am assuming its not solvable at the moment, or I misunderstood your reply.

I will be closing issue in a few days

from transformers-neuronx.