+- Inference (https://www.datasculpting.ai/inference)
+-- Forum: Drivers / Frameworks / Toolkits (https://www.datasculpting.ai/inference/forumdisplay.php?fid=8)
+--- Forum: PyTorch / JAX / Other Frameworks (https://www.datasculpting.ai/inference/forumdisplay.php?fid=34)
+--- Thread: JAX for inference only on 4090 (/showthread.php?tid=494)
Pages:
1
2
JAX for inference only on 4090 - torchstation - 09-26-2026
This keeps coming up in conversation so I figured I'd start a thread.
One thing I keep noticing is that people underestimate tail latency until they actually try to run this stuff for more than one afternoon.
Curious whether that lines up with what other people are seeing.
This is mostly for local work rather than public serving.
I expected the biggest difference to come from the model, but the stack around it mattered at least as much.
Re: JAX for inference only on 4090 - southspark - 09-27-2026
I think the bottleneck is somewhere else. The reported numbers can look compelling, but in ordinary usage I didn't see the same benefit.
Re: JAX for inference only on 4090 - southden - 09-27-2026
Interesting, I hadn't tried that. I don't think your conclusion is far off. The biggest change for me was simplifying the stack instead of adding one more optimization layer.
Re: JAX for inference only on 4090 - plainscope50 - 09-29-2026
A couple things I'd check first. Before rebuilding anything, I'd test one clean baseline with Phi and compare from there. Once you know the plain path works, then add the faster path back in piece by piece.
I'm happy to test one more variant if you want a second data point.
Re: JAX for inference only on 4090 - modgarden - 09-29-2026
Quote:A couple things I'd check first. Before rebuilding anything, I'd test one clean baseline with Phi and compare from there. Once you know the plain path
What worked for me was pretty simple. On B200, the easiest win is often to back off the aggressive settings and see whether the problem survives. Once you know the plain path works, then add the faster path back in piece by piece.
Re: JAX for inference only on 4090 - emberrunner - 09-29-2026
Not sure I agree there. I've seen people blame INT4 when the real issue was much more mundane. Could be a workload difference, but I wouldn't generalize too fast from one setup.
Re: JAX for inference only on 4090 - llmarch51 - 09-30-2026
Appreciate the detail. I don't think your conclusion is far off.
Re: JAX for inference only on 4090 - oxideworks59 - 09-30-2026
Need a little more detail here. Is this local-only or are you routing anything externally?
Re: JAX for inference only on 4090 - mangopeak72 - 09-30-2026
What worked for me was pretty simple. If you're on llama.cpp, I'd confirm the version first and then reduce variables until the behavior is obvious.
Re: JAX for inference only on 4090 - turboscope - 09-30-2026
I tested this tonight. On A6000 with MLX, I saw basically the same behavior. I expected the model choice to matter more than it did.