Roll back small number of KV cache updates

[mseeger]

Is your feature request related to a problem? Please describe.
I'd like to be able to roll back a small number of KV cache updates, without having to checkpoint the full cache. Arises in different situations:

• Evaluation: Compute different metrics, given the same prompt. Or sample different responses
• Speculative decoding

Describe the solution you'd like
This is essentially just a special case of the annotations used in autograd_hooks.py. We need to store index and delta for every update, so that cache buffers can be restored. Ideally, this is supposed by some methods in the KVCache base class.

[mseeger]

Better solution: We also need a mechanism to switch off CPU cache offloading during token generation:

• Cache can switch in a new buffer setup, with content copied from the existing one. The old one is kept (may be offloaded)
• Token are generated with new buffer setup
• At end: Old buffers are switched in again

This is much simpler, solves two problems, and rolling back could interfere with buffer quantization.

Branch: switch_buffers

[mseeger]

WAIT! Some issues:

• CPU offloading quite slow at the moment: Try not to use it
• Can't use non-quant buffers during generation: Not enough memory
• Buffer switching: Cannot have 2 sets of buffers in memory!
• Quantized buffers during generation: Need to quantize and write back after every token

Better idea:

• Process prompt with quantized buffers, as usual
• For generation (token by token): Switch to "replay buffers"

Replay buffers: Stores link to quantized buffers and list of scatter updates

For each token and layer:

• Recreate buffer content by first de-quantize, then reply scatter updates
• Do the single token update
• Just append the new scatter update

This is very elegant: Don't have to checkpoint the buffers at all, because the updates during token generation are never written back.

Replay wrapper for quantized cache buffers:

• Wraps list of QuantizedKVCacheBuffers, connected to common DequantizedKVCacheBuffers
• Result is list of ReplayKVCacheBuffers, connected to the same DequantizedKVCacheBuffers. Each object links to its QuantizedKVCacheBuffers and stores scatter update info as a list
• Supports updates of size q_len=1 only
• Support converting back to QuantizedKVCacheBuffers. This would drive tool calling

[mseeger]

Had to roll back #89 commit for now: Obtained lower metric values than before.

Branch: switch_buffers_revisit.

Get back to this one, this is important, but have to get it right, including a test comparing logits:

• Old way: model(prompts, targets=None), then generate
• New way: As done here

[mseeger]

Branch: switch_buffers4.

On hold again:

• test_replay_buffers: Still two cases failing, on logits for first token already
• val_time_avg5 / val_time_avg1 = 1.77 (+- 0.167): Can average score-based metrics 5x. But this does not make a big difference over the validation set