add support for batched input_pos to model

[t-vi]

This allows to feed both x, input_pos of shape (batch, T) to the model.
I included a test that compares to running the batch items separately.

To add some commentary around the for loop:

• it's not terribly cool to use the for loop, but to reduce to index_copy_, we would need to fold the batch and indexed dimensions, which we cannot for KV-Cache,
• if we can use cudagraphs, it might not be too bad.