Move generation functionality to base classes

[mattdangerw]

We want it to be easy to add more generative models to the repository, and we want it to be easy to add new features related to generation.

Currently quite a bit of common functionality for generation live in every XXCausalLM subclass, which makes it hard to add new models and build new generation features. I think we can pull a lot of this functionality down into the base class. I also as part of this work we should move the compiled while loop out of the sampler class and into the CausalLM base class, which will make it much easier to write custom generative loops (e.g. would help integrating with JetStream).

[mattdangerw]

This is mostly implemented, but still needs a little work. I'll push code shortly.