enable KV cache in miniGPT

[windmaple]

This PR enables KV cache for miniGPT inference (inference time on goes from ~9s to ~3s on my Cloudtop). Also updated to use simpler sharding annotation with newer NNX API.

Correctness is validated in https://colab.research.google.com/drive/1Fw2IQjH-UcGReOXw6ykqJaXKWUv3HN_O?resourcekey=0-lNpYdIeKUxoOMfG_KpZfAw&usp=sharing

[windmaple]

@emilyfertig This is a bit beyond the original revamping scope but I think it's nice to have because 1) we need to use the newer sharding annotation anyway at some point 2) KV cache is very commonly used for LLM inference. This would be a good reference regardless of whether we can get vLLM TPU integration.

[salfaris]

I was just about to suggest the updated nnx.with_partitioning changes! This is great work :)

[emilyfertig]

Thanks!