Supporting varlen?

[sh0416]

Is there any plan to support varlen interface? It seems useful for batch generation and some kv cache prefetch scheme would enhance the performance.

[nshepperd]

Varlen is partially supported now. It fully works, except there is no vmap handling yet, and no support for automatic partitioning of it in multi-gpu settings yet (so you'd have to use shard_map by hand in this case). I'll add that... when i get around to it.