Refactoring of `GPT.forward` when it comes to `input_pos` and KV cache usage

The current GPT.forward in model.py essentially serves two use cases:

• Forward pass for training: input_pos=None, KV cache not used. Implicitly, input_pos = arange(idx.shape[-1]), and causal masking is used. Could also be used for prefill with prompt in inference.
• Inference. input_pos is not None, KV cache is used. There seem two cases here, either input_pos = arange(idx.shape[-1]) (used for prefill), or input_pos.shape[-1] == 1 (generation of single next token, possibly batched).

I am interested in implementing KV cache strategies, such as H2O. In inference, we really only have prefill, and then single-token generation. Inference always works like this:

• Prefill with sequence length T (minimum of prompt size and max cache size)
• Generate token T
• Generate token T+1
• ...

Most KV cache strategies only support this protocol.

My proposal would be to refactor GPT.forward to support two cases only:

• Forward pass for training: With an additional flag, this can be used for prefill, in that this would initialize the KV cache with the K and V vectors obtained as part of the forward, just because scaled_dot_product_attention is called
• Generate single token: idx.shape[-1] == 1. input_pos is not really needed, it would rather be input_pos_maxp1. The KV cache tracks the position of the next token, and it would complain if asked to do anything else

This supports everything you have right now, plus it supports advanced KV caches like H2O.

I am happy to do this in a branch in my fork and show you how it would look like.

Start by reading model.py and the current GPT.forward implementation, focusing on how input_pos, causal masking, and the KV cache are handled. Compare the training/prefill and single-token generation paths against the proposed two-case protocol. Done should preserve current behavior while allowing KV cache strategies such as H2O.