Add KV cache

[AnMakc]

Added optional KV cache to speed up inference.

Cache only added for main transformer layers, not s2 inference (DependencyAwareLayer). Not sure if adding it make sense as there is only one attention block.

Cache resets every step after (context len + num generated tokens ) exceeds max_context. This does not affect outputs, but performance will degrade to non-cached version. To avoid cache reset, keep (context + pred_len) under max_context (e.g. 480 + 32).
Cache rolling resulted in regression in one test so I just kept things safe. It's possible that this is a flaky test, but I was not able to debug it.

Results

On A100 with 400 context throughput improved 5.17 -> 44.3 batch/sec (165 -> 1418 token/sec). Saturates with ~256 batch size at 6.8 batch/sec (1740 token/sec).

[AnMakc]

@shiyu-coder For the cache rolling not working - it seems to be effect of the RoPE embeddings being relative to window start, so when the context window slides, embeddings effectively change for all observed tokens. If you have ideas on how to fix this - cache resets could be avoided. But I wouldn't add this to PR as current version works for most cases and produces stable results.