Rewind/Restore/Fork points into the state

[cheery]

I'm planning to make a chat session where you can rewind the discussion and continue from there. I was planning of using save states for this, but saving the whole state every time is very expensive.

I'm not entirely clear on how the state works. Would it be possible to piece the state itself such that it is rewindable and that I would only have to store one state per session?

[cheery]

If I were to guess what I have to do...

• First store seq with llama_state_seq_get_data.
• then retrieve logits.
• retrieve logits every time you want a fork point.
• when it's time to run from a fork, load the seq, truncate it, then feed old logits to the sampler.

If I'm not correct, I would appreciate correction.

Updates: I did learn about KV caches and attention masks. I maybe need to store the whole state here.

Logits can be stored by giving a decoder batch where it has been instructed to save logits. Then they're in state image and I can truncate the sequence to previous logit boundary. This is what I'm going to try eventually.

[cheery]

I was thinking on this and realising this kind of a chat log forms a tree structure. Each call to generate or user appended message is a new leaf to the tree.

It would really help me out if I knew more about the structure of context.

1. When I call decode(), what exactly happens?
2. What data needs to be stored in order to continue generation without having to call decode() again?
3. Can that data be pieced up into smaller chunks?