JAX for inference only

[tinycontext]

Posting this in case it saves someone else a little time.

I'm helping a small team set up a local workflow. We don't need a giant stack, but we do care about prefill/decode split.

Would rather learn from your scars than collect my own all at once.

[torchsupply99]

Need a little more detail here. Are you testing this on GB200 or something else?

[cobalttinker]

Good point. I don't think your conclusion is far off. I got better results once I stopped chasing one metric and just looked at whether the workflow felt usable.